Merge branch 'pb-fixed-mocked-ci-connection-2' into 'master'

Only mock ci connections in single db mode

See merge request gitlab-org/gitlab!74846

spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb

@@ -487,25 +487,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
   end
 
-  describe 'primary connection re-use', :reestablished_active_record_base do
+  describe 'primary connection re-use', :reestablished_active_record_base, :add_ci_connection do
     let(:model) { Ci::ApplicationRecord }
 
-    around do |example|
-      if Gitlab::Database.has_config?(:ci)
-        example.run
-      else
-        # fake additional Database
-        model.establish_connection(
-          ActiveRecord::DatabaseConfigurations::HashConfig.new(Rails.env, 'ci', ActiveRecord::Base.connection_db_config.configuration_hash)
-        )
-
-        example.run
-
-        # Cleanup connection_specification_name for Ci::ApplicationRecord
-        model.remove_connection
-      end
-    end
-
     describe '#read' do
       it 'returns ci replica connection' do
         expect { |b| lb.read(&b) }.to yield_with_args do |args|

spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb

@@ -17,7 +17,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics, query_ana
     process_sql(ActiveRecord::Base, "SELECT 1 FROM projects")
   end
 
-  context 'properly observes all queries', :mocked_ci_connection do
+  context 'properly observes all queries', :add_ci_connection do
     using RSpec::Parameterized::TableSyntax
 
     where do

spec/support/database/multiple_databases.rb

@@ -6,6 +6,10 @@ module Database
       skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
     end
 
+    def skip_if_multiple_databases_are_setup
+      skip 'Skipping because multiple databases are set up' if Gitlab::Database.has_config?(:ci)
+    end
+
     def reconfigure_db_connection(name: nil, config_hash: {}, model: ActiveRecord::Base, config_model: nil)
       db_config = (config_model || model).connection_db_config
 
@@ -46,6 +50,26 @@ module Database
       new_handler&.clear_all_connections!
     end
     # rubocop:enable Database/MultipleDatabases
+
+    def with_added_ci_connection
+      if Gitlab::Database.has_config?(:ci)
+        # No need to add a ci: connection if we already have one
+        yield
+      else
+        with_reestablished_active_record_base(reconnect: true) do
+          reconfigure_db_connection(
+            name: :ci,
+            model: Ci::ApplicationRecord,
+            config_model: ActiveRecord::Base
+          )
+
+          yield
+
+          # Cleanup connection_specification_name for Ci::ApplicationRecord
+          Ci::ApplicationRecord.remove_connection
+        end
+      end
+    end
   end
 
   module ActiveRecordBaseEstablishConnection
@@ -69,18 +93,9 @@ RSpec.configure do |config|
     end
   end
 
-  config.around(:each, :mocked_ci_connection) do |example|
-    with_reestablished_active_record_base(reconnect: true) do
-      reconfigure_db_connection(
-        name: :ci,
-        model: Ci::ApplicationRecord,
-        config_model: ActiveRecord::Base
-      )
-
+  config.around(:each, :add_ci_connection) do |example|
+    with_added_ci_connection do
       example.run
-
-      # Cleanup connection_specification_name for Ci::ApplicationRecord
-      Ci::ApplicationRecord.remove_connection
     end
   end
 end

spec/support_specs/database/multiple_databases_spec.rb

@@ -56,4 +56,43 @@ RSpec.describe 'Database::MultipleDatabases' do
       end
     end
   end
+
+  describe '.with_added_ci_connection' do
+    context 'when only a single database is setup' do
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      it 'connects Ci::ApplicationRecord to the main database for the duration of the block', :aggregate_failures do
+        main_database = current_database(ActiveRecord::Base)
+        original_database = current_database(Ci::ApplicationRecord)
+
+        with_added_ci_connection do
+          expect(current_database(Ci::ApplicationRecord)).to eq(main_database)
+        end
+
+        expect(current_database(Ci::ApplicationRecord)).to eq(original_database)
+      end
+    end
+
+    context 'when multiple databases are setup' do
+      before do
+        skip_if_multiple_databases_not_setup
+      end
+
+      it 'does not mock the original Ci::ApplicationRecord connection', :aggregate_failures do
+        original_database = current_database(Ci::ApplicationRecord)
+
+        with_added_ci_connection do
+          expect(current_database(Ci::ApplicationRecord)).to eq(original_database)
+        end
+
+        expect(current_database(Ci::ApplicationRecord)).to eq(original_database)
+      end
+    end
+
+    def current_database(connection_class)
+      connection_class.retrieve_connection.execute('select current_database()').first
+    end
+  end
 end