Merge branch 'pb-fix-setup-with-shared-connections' into 'master'

Fix gitlab:setup when connections are shared

See merge request gitlab-org/gitlab!83319

lib/gitlab/database/each_database.rb

@@ -4,11 +4,13 @@ module Gitlab
   module Database
     module EachDatabase
       class << self
-        def each_database_connection(only: nil)
+        def each_database_connection(only: nil, include_shared: true)
           selected_names = Array.wrap(only)
           base_models = select_base_models(selected_names)
 
           base_models.each_pair do |connection_name, model|
+            next if !include_shared && Gitlab::Database.db_config_share_with(model.connection_db_config)
+
             connection = model.connection
 
             with_shared_connection(connection, connection_name) do

lib/tasks/gitlab/setup.rake

@@ -53,7 +53,7 @@ namespace :gitlab do
         AND pid <> pg_backend_pid();
     SQL
 
-    Gitlab::Database::EachDatabase.each_database_connection do |connection|
+    Gitlab::Database::EachDatabase.each_database_connection(include_shared: false) do |connection|
       connection.execute(cmd)
     rescue ActiveRecord::NoDatabaseError
     end

spec/lib/gitlab/database/each_database_spec.rb

@@ -58,6 +58,15 @@ RSpec.describe Gitlab::Database::EachDatabase do
         end
       end
     end
+
+    context 'when shared connections are not included' do
+      it 'only yields the unshared connections' do
+        expect(Gitlab::Database).to receive(:db_config_share_with).twice.and_return(nil, 'main')
+
+        expect { |b| described_class.each_database_connection(include_shared: false, &b) }
+          .to yield_successive_args([ActiveRecord::Base.connection, 'main'])
+      end
+    end
   end
 
   describe '.each_model_connection' do

spec/tasks/gitlab/setup_rake_spec.rb

@@ -22,7 +22,11 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
     let(:server_service1) { double(:server_service) }
     let(:server_service2) { double(:server_service) }
 
-    let(:connections) { Gitlab::Database.database_base_models.values.map(&:connection) }
+    let(:connections) do
+      Gitlab::Database.database_base_models.values.filter_map do |model|
+        model.connection if Gitlab::Database.db_config_share_with(model.connection_db_config).nil?
+      end
+    end
 
     before do
       allow(Gitlab).to receive_message_chain('config.repositories.storages').and_return(storages)
@@ -119,6 +123,10 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
     end
 
     def expect_connections_to_be_terminated
+      expect(Gitlab::Database::EachDatabase).to receive(:each_database_connection)
+        .with(include_shared: false)
+        .and_call_original
+
       expect(connections).to all(receive(:execute).with(/SELECT pg_terminate_backend/))
     end