1 files changed, 18 insertions, 15 deletions
diff --git a/lib/initializer_connections.rb b/lib/initializer_connections.rb
index c8a6bb6c511..ae2809b7604 100644
--- a/lib/initializer_connections.rb
+++ b/lib/initializer_connections.rb
@@ -1,29 +1,32 @@
 # frozen_string_literal: true
 
 module InitializerConnections
-  # Prevents any database connections within the block
-  # by using an empty connection handler
-  # rubocop:disable Database/MultipleDatabases
-  def self.with_disabled_database_connections
+  # Raises if new database connections established within the block
+  #
+  # NOTE: this does not prevent existing connections that is already checked out
+  # from being used. You will need other means to prevent that such as by
+  # clearing all connections as implemented in the
+  # `:clear_active_connections_again` initializer for routes
+  #
+  def self.raise_if_new_database_connection
     return yield if Gitlab::Utils.to_boolean(ENV['SKIP_RAISE_ON_INITIALIZE_CONNECTIONS'])
 
-    original_handler = ActiveRecord::Base.connection_handler
-
-    dummy_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
-    ActiveRecord::Base.connection_handler = dummy_handler
+    previous_connection_counts = ActiveRecord::Base.connection_handler.connection_pool_list.to_h do |pool|
+      [pool.db_config.name, pool.connections.size]
+    end
 
     yield
 
-    if dummy_handler&.connection_pool_names&.present?
-      raise "Unxpected connection_pools (#{dummy_handler.connection_pool_names}) ! Call `connects_to` before this block"
+    new_connection_counts = ActiveRecord::Base.connection_handler.connection_pool_list.to_h do |pool|
+      [pool.db_config.name, pool.connections.size]
     end
-  rescue ActiveRecord::ConnectionNotEstablished
+
+    raise_database_connection_made_error unless previous_connection_counts == new_connection_counts
+  end
+
+  def self.raise_database_connection_made_error
     message = "Database connection should not be called during initializers. Read more at https://docs.gitlab.com/ee/development/rails_initializers.html#database-connections-in-initializers"
 
     raise message
-  ensure
-    ActiveRecord::Base.connection_handler = original_handler
-    dummy_handler&.clear_all_connections!
   end
-  # rubocop:enable Database/MultipleDatabases
 end