Add `#with_reestablished_active_record_base` to allow usage of `establish_connection`

ee/spec/lib/gitlab/geo/database_tasks_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Geo::DatabaseTasks do
+RSpec.describe Gitlab::Geo::DatabaseTasks, :reestablished_active_record_base do
   let(:schema_file) { Rails.root.join('tmp', 'tests', 'geo_schema.rb').to_s }
 
   subject { described_class }

ee/spec/lib/system_check/geo/geo_database_configured_check_spec.rb

@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
   subject { described_class.new }
 
-  describe '#multi_check' do
+  describe '#multi_check', :reestablished_active_record_base do
     it "checks database configuration" do
       stub_configuration_check(false)
 

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -235,7 +235,7 @@ module Gitlab
             @configuration.model.connection_specification_name,
             role: ActiveRecord::Base.writing_role,
             shard: ActiveRecord::Base.default_shard
-          )
+          ) || raise(::ActiveRecord::ConnectionNotEstablished)
         end
 
         private

spec/initializers/database_config_spec.rb

@@ -2,19 +2,11 @@
 
 require 'spec_helper'
 
-RSpec.describe 'Database config initializer' do
+RSpec.describe 'Database config initializer', :reestablished_active_record_base do
   subject do
     load Rails.root.join('config/initializers/database_config.rb')
   end
 
-  around do |example|
-    original_config = ActiveRecord::Base.connection_db_config
-
-    example.run
-
-    ActiveRecord::Base.establish_connection(original_config)
-  end
-
   before do
     allow(Gitlab::Runtime).to receive(:max_threads).and_return(max_threads)
   end

spec/lib/gitlab/database/bulk_update_spec.rb

@@ -91,45 +91,38 @@ RSpec.describe Gitlab::Database::BulkUpdate do
       .to eq(['MR a', 'Issue a', 'Issue b'])
   end
 
-  shared_examples 'basic functionality' do
-    it 'sets multiple values' do
-      create_default(:user)
-      create_default(:project)
-
-      i_a, i_b = create_list(:issue, 2)
+  context 'validates prepared_statements support', :reestablished_active_record_base do
+    using RSpec::Parameterized::TableSyntax
 
-      mapping = {
-        i_a => { title: 'Issue a' },
-        i_b => { title: 'Issue b' }
-      }
+    where(:prepared_statements) do
+      [false, true]
+    end
 
-      described_class.execute(%i[title], mapping)
+    before do
+      configuration_hash = ActiveRecord::Base.connection_db_config.configuration_hash
 
-      expect([i_a, i_b].map { |x| x.reset.title })
-        .to eq(['Issue a', 'Issue b'])
+      ActiveRecord::Base.establish_connection(
+        configuration_hash.merge(prepared_statements: prepared_statements)
+      )
     end
-  end
 
-  include_examples 'basic functionality'
+    with_them do
+      it 'sets multiple values' do
+        create_default(:user)
+        create_default(:project)
 
-  context 'when prepared statements are configured differently to the normal test environment' do
-    before do
-      klass = Class.new(ActiveRecord::Base) do
-        def self.abstract_class?
-          true # So it gets its own connection
-        end
-      end
+        i_a, i_b = create_list(:issue, 2)
 
-      stub_const('ActiveRecordBasePreparedStatementsInverted', klass)
+        mapping = {
+          i_a => { title: 'Issue a' },
+          i_b => { title: 'Issue b' }
+        }
 
-      c = ActiveRecord::Base.retrieve_connection.instance_variable_get(:@config)
-      inverted = c.merge(prepared_statements: !ActiveRecord::Base.connection.prepared_statements)
-      ActiveRecordBasePreparedStatementsInverted.establish_connection(inverted)
+        described_class.execute(%i[title], mapping)
 
-      allow(ActiveRecord::Base).to receive(:connection_specification_name)
-        .and_return(ActiveRecordBasePreparedStatementsInverted.connection_specification_name)
+        expect([i_a, i_b].map { |x| x.reset.title })
+          .to eq(['Issue a', 'Issue b'])
+      end
     end
-
-    include_examples 'basic functionality'
   end
 end

spec/lib/gitlab/database/connection_spec.rb

@@ -126,15 +126,7 @@ RSpec.describe Gitlab::Database::Connection do
     end
   end
 
-  describe '#disable_prepared_statements' do
-    around do |example|
-      original_config = connection.scope.connection.pool.db_config
-
-      example.run
-
-      connection.scope.establish_connection(original_config)
-    end
-
+  describe '#disable_prepared_statements', :reestablished_active_record_base do
     it 'disables prepared statements' do
       connection.scope.establish_connection(
         ::Gitlab::Database.main.config.merge(prepared_statements: true)

spec/lib/gitlab/database/postgresql_adapter/force_disconnectable_mixin_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin do
+RSpec.describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin, :reestablished_active_record_base do
   describe 'checking in a connection to the pool' do
     let(:model) do
       Class.new(ActiveRecord::Base) do

spec/lib/gitlab/database/schema_migrations/context_spec.rb

@@ -23,7 +23,7 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
       end
     end
 
-    context 'multiple databases' do
+    context 'multiple databases', :reestablished_active_record_base do
       let(:connection_class) do
         Class.new(::ApplicationRecord) do
           self.abstract_class = true
@@ -34,8 +34,6 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
         end
       end
 
-      let(:configuration_overrides) { {} }
-
       before do
         connection_class.establish_connection(
           ActiveRecord::Base

spec/support/database/multiple_databases.rb

@@ -5,5 +5,57 @@ module Database
     def skip_if_multiple_databases_not_setup
       skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
     end
+
+    # The usage of this method switches temporarily used `connection_handler`
+    # allowing full manipulation of ActiveRecord::Base connections without
+    # having side effects like:
+    # - misaligned transactions since this is managed by `BeforeAllAdapter`
+    # - removal of primary connections
+    #
+    # The execution within a block ensures safe cleanup of all allocated resources.
+    #
+    # rubocop:disable Database/MultipleDatabases
+    def with_reestablished_active_record_base(reconnect: true)
+      connection_classes = ActiveRecord::Base.connection_handler.connection_pool_names.map(&:constantize).to_h do |klass|
+        [klass, klass.connection_db_config]
+      end
+
+      original_handler = ActiveRecord::Base.connection_handler
+      new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+      ActiveRecord::Base.connection_handler = new_handler
+
+      if reconnect
+        connection_classes.each { |klass, db_config| klass.establish_connection(db_config) }
+      end
+
+      yield
+    ensure
+      ActiveRecord::Base.connection_handler = original_handler
+      new_handler&.clear_all_connections!
+    end
+    # rubocop:enable Database/MultipleDatabases
+  end
+
+  module ActiveRecordBaseEstablishConnection
+    def establish_connection(*args)
+      # rubocop:disable Database/MultipleDatabases
+      if connected? && connection&.transaction_open? && ActiveRecord::Base.connection_handler == ActiveRecord::Base.default_connection_handler
+        raise "Cannot re-establish '#{self}.establish_connection' within an open transaction (#{connection&.open_transactions.to_i}). " \
+          "Use `with_reestablished_active_record_base` instead or add `:reestablished_active_record_base` to rspec context."
+      end
+      # rubocop:enable Database/MultipleDatabases
+
+      super
+    end
   end
 end
+
+RSpec.configure do |config|
+  config.around(:each, :reestablished_active_record_base) do |example|
+    with_reestablished_active_record_base(reconnect: example.metadata.fetch(:reconnect, true)) do
+      example.run
+    end
+  end
+end
+
+ActiveRecord::Base.singleton_class.prepend(::Database::ActiveRecordBaseEstablishConnection) # rubocop:disable Database/MultipleDatabases

spec/support_specs/database/multiple_databases_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Database::MultipleDatabases' do
+  describe '.with_reestablished_active_record_base' do
+    context 'when doing establish_connection' do
+      context 'on ActiveRecord::Base' do
+        it 'raises exception' do
+          expect { ActiveRecord::Base.establish_connection(:main) }.to raise_error /Cannot re-establish/
+        end
+
+        context 'when using with_reestablished_active_record_base' do
+          it 'does not raise exception' do
+            with_reestablished_active_record_base do
+              expect { ActiveRecord::Base.establish_connection(:main) }.not_to raise_error
+            end
+          end
+        end
+      end
+
+      context 'on Ci::CiDatabaseRecord' do
+        before do
+          skip_if_multiple_databases_not_setup
+        end
+
+        it 'raises exception' do
+          expect { Ci::CiDatabaseRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
+        end
+
+        context 'when using with_reestablished_active_record_base' do
+          it 'does not raise exception' do
+            with_reestablished_active_record_base do
+              expect { Ci::CiDatabaseRecord.establish_connection(:main) }.not_to raise_error
+            end
+          end
+        end
+      end
+    end
+
+    context 'when trying to access connection' do
+      context 'when reconnect is true' do
+        it 'does not raise exception' do
+          with_reestablished_active_record_base(reconnect: true) do
+            expect { ActiveRecord::Base.connection.execute("SELECT 1") }.not_to raise_error # rubocop:disable Database/MultipleDatabases
+          end
+        end
+      end
+
+      context 'when reconnect is false' do
+        it 'does raise exception' do
+          with_reestablished_active_record_base(reconnect: false) do
+            expect { ActiveRecord::Base.connection.execute("SELECT 1") }.to raise_error(ActiveRecord::ConnectionNotEstablished) # rubocop:disable Database/MultipleDatabases
+          end
+        end
+      end
+    end
+  end
+end