Allow nesting of SharedModel.using_connection

Allow SharedModel.using_connection to be called within nested blocks
with the same connection, for situations where we want to wrap multiple
entry points that might be called in a chain.

lib/gitlab/database/shared_model.rb

@@ -8,13 +8,17 @@ module Gitlab
 
       class << self
         def using_connection(connection)
-          raise 'cannot nest connection overrides for shared models' unless overriding_connection.nil?
+          previous_connection = self.overriding_connection
+
+          unless previous_connection.nil? || previous_connection.equal?(connection)
+            raise 'cannot nest connection overrides for shared models with different connections'
+          end
 
           self.overriding_connection = connection
 
           yield
         ensure
-          self.overriding_connection = nil
+          self.overriding_connection = nil unless previous_connection.equal?(self.overriding_connection)
         end
 
         def connection

spec/lib/gitlab/database/shared_model_spec.rb

@@ -27,6 +27,38 @@ RSpec.describe Gitlab::Database::SharedModel do
       end
     end
 
+    context 'when multiple connection overrides are nested', :aggregate_failures do
+      let(:second_connection) { double('connection') }
+
+      it 'allows the nesting with the same connection object' do
+        expect_original_connection_around do
+          described_class.using_connection(new_connection) do
+            expect(described_class.connection).to be(new_connection)
+
+            described_class.using_connection(new_connection) do
+              expect(described_class.connection).to be(new_connection)
+            end
+
+            expect(described_class.connection).to be(new_connection)
+          end
+        end
+      end
+
+      it 'raises an error if the connection is changed' do
+        expect_original_connection_around do
+          described_class.using_connection(new_connection) do
+            expect(described_class.connection).to be(new_connection)
+
+            expect do
+              described_class.using_connection(second_connection) {}
+            end.to raise_error(/cannot nest connection overrides/)
+
+            expect(described_class.connection).to be(new_connection)
+          end
+        end
+      end
+    end
+
     context 'when the block raises an error', :aggregate_failures do
       it 're-raises the error, removing the overridden connection' do
         expect_original_connection_around do