Merge branch 'rework_with_disabled_database_connections_for_rails_7' into 'master'

Rework with_disabled_database_connections to work with Rails 7.0

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/111399



Merged-by: Ethan Urie <eurie@gitlab.com>
Approved-by: Rutger Wessels <rwessels@gitlab.com>
Approved-by: Ethan Urie <eurie@gitlab.com>
Reviewed-by: Rutger Wessels <rwessels@gitlab.com>
Co-authored-by: Thong Kuah <tkuah@gitlab.com>

config/initializers/database_config.rb

@@ -2,14 +2,14 @@
 Rails.application.reloader.to_run(:before) do
   # Make sure connects_to for Ci::ApplicationRecord gets called outside of config/routes.rb first
-  # See InitializerConnections.with_disabled_database_connections
+  # See InitializerConnections.raise_if_new_database_connection
   Ci::ApplicationRecord
 end
 Gitlab.ee do
   if Gitlab::Geo.geo_database_configured?
     # Make sure connects_to for geo gets called outside of config/routes.rb first
-    # See InitializerConnections.with_disabled_database_connections
+    # See InitializerConnections.raise_if_new_database_connection
     Geo::TrackingBase
   end

config/routes.rb

@@ -3,7 +3,7 @@
 require 'sidekiq/web'
 require 'sidekiq/cron/web'
-InitializerConnections.with_disabled_database_connections do
+InitializerConnections.raise_if_new_database_connection do
   Rails.application.routes.draw do
     concern :access_requestable do
       post :request_access, on: :collection

lib/initializer_connections.rb

 # 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

rubocop/cop/database/multiple_databases.rb

@@ -20,6 +20,7 @@ class MultipleDatabases < RuboCop::Cop::Base
         ALLOWED_METHODS = %i[
           no_touching
           configurations
+          connection_handler
           logger
         ].freeze

spec/config/application_spec.rb

@@ -3,31 +3,57 @@
 require 'spec_helper'
 RSpec.describe Gitlab::Application, feature_category: :scalability do # rubocop:disable RSpec/FilePath
-  using RSpec::Parameterized::TableSyntax
-  filtered_param = ActiveSupport::ParameterFilter::FILTERED
-  context 'when parameters are logged' do
-    describe 'rails does not leak confidential parameters' do
-      def request_for_url(input_url)
-        env = Rack::MockRequest.env_for(input_url)
-        env['action_dispatch.parameter_filter'] = described_class.config.filter_parameters
-        ActionDispatch::Request.new(env)
-      end
-      where(:input_url, :output_query) do
-        '/'                                      | {}
-        '/?safe=1'                               | { 'safe' => '1' }
-        '/?private_token=secret'                 | { 'private_token' => filtered_param }
-        '/?mixed=1&private_token=secret'         | { 'mixed' => '1', 'private_token' => filtered_param }
-        '/?note=secret&noteable=1&prefix_note=2' | { 'note' => filtered_param, 'noteable' => '1', 'prefix_note' => '2' }
-        '/?note[note]=secret&target_type=1'      | { 'note' => filtered_param, 'target_type' => '1' }
-        '/?safe[note]=secret&target_type=1'      | { 'safe' => { 'note' => filtered_param }, 'target_type' => '1' }
+  describe 'config.filter_parameters' do
+    using RSpec::Parameterized::TableSyntax
+    filtered = ActiveSupport::ParameterFilter::FILTERED
+    context 'when parameters are logged' do
+      describe 'rails does not leak confidential parameters' do
+        def request_for_url(input_url)
+          env = Rack::MockRequest.env_for(input_url)
+          env['action_dispatch.parameter_filter'] = described_class.config.filter_parameters
+          ActionDispatch::Request.new(env)
+        end
+
+        where(:input_url, :output_query) do
+          '/'                                      | {}
+          '/?safe=1'                               | { 'safe' => '1' }
+          '/?private_token=secret'                 | { 'private_token' => filtered }
+          '/?mixed=1&private_token=secret'         | { 'mixed' => '1', 'private_token' => filtered }
+          '/?note=secret&noteable=1&prefix_note=2' | { 'note' => filtered, 'noteable' => '1', 'prefix_note' => '2' }
+          '/?note[note]=secret&target_type=1'      | { 'note' => filtered, 'target_type' => '1' }
+          '/?safe[note]=secret&target_type=1'      | { 'safe' => { 'note' => filtered }, 'target_type' => '1' }
+        end
+        with_them do
+          it { expect(request_for_url(input_url).filtered_parameters).to eq(output_query) }
+        end
       end
-      with_them do
-        it { expect(request_for_url(input_url).filtered_parameters).to eq(output_query) }
+    end
+  end
+
+  describe 'clear_active_connections_again initializer' do
+    subject(:clear_active_connections_again) do
+      described_class.initializers.find { |i| i.name == :clear_active_connections_again }
+    end
+
+    it 'is included in list of Rails initializers' do
+      expect(clear_active_connections_again).to be_present
+    end
+
+    it 'is configured after set_routes_reloader_hook' do
+      expect(clear_active_connections_again.after).to eq(:set_routes_reloader_hook)
+    end
+
+    describe 'functionality', :reestablished_active_record_base do
+      it 'clears all connections' do
+        Project.first
+
+        clear_active_connections_again.run
+        expect(ActiveRecord::Base.connection_handler.active_connections?).to eq(false)
       end
     end
   end

spec/lib/initializer_connections_spec.rb

@@ -3,15 +3,20 @@
 require 'spec_helper'
 RSpec.describe InitializerConnections do
-  describe '.with_disabled_database_connections', :reestablished_active_record_base do
+  describe '.raise_if_new_database_connection', :reestablished_active_record_base do
+    before do
+      ActiveRecord::Base.connection_handler.clear_active_connections!
+      ActiveRecord::Base.connection_handler.flush_idle_connections!
+    end
+
     def block_with_database_call
-      described_class.with_disabled_database_connections do
+      described_class.raise_if_new_database_connection do
         Project.first
       end
     end
     def block_with_error
-      described_class.with_disabled_database_connections do
+      described_class.raise_if_new_database_connection do
         raise "oops, an error"
       end
     end
@@ -20,6 +25,12 @@ def block_with_error
       expect { block_with_database_call }.to raise_error(/Database connection should not be called during initializer/)
     end
+    it 'prevents any database connection re-use within the block' do
+      Project.connection # establish a connection first, it will be used inside the block
+
+      expect { block_with_database_call }.to raise_error(/Database connection should not be called during initializer/)
+    end
+
     it 'does not prevent database connection if SKIP_RAISE_ON_INITIALIZE_CONNECTIONS is set' do
       stub_env('SKIP_RAISE_ON_INITIALIZE_CONNECTIONS', '1')
@@ -33,31 +44,34 @@ def block_with_error
     end
     it 'restores original connection handler' do
-      # rubocop:disable Database/MultipleDatabases
       original_handler = ActiveRecord::Base.connection_handler
       expect { block_with_database_call }.to raise_error(/Database connection should not be called during initializer/)
       expect(ActiveRecord::Base.connection_handler).to eq(original_handler)
-      # rubocop:enabled Database/MultipleDatabases
     end
     it 'restores original connection handler even there is an error' do
-      # rubocop:disable Database/MultipleDatabases
       original_handler = ActiveRecord::Base.connection_handler
       expect { block_with_error }.to raise_error(/an error/)
       expect(ActiveRecord::Base.connection_handler).to eq(original_handler)
-      # rubocop:enabled Database/MultipleDatabases
     end
-    it 'raises if any new connection_pools are established in the block' do
+    it 'does not raise if connection_pool is retrieved in the block' do
+      # connection_pool, connection_db_config doesn't connect to database, so it's OK
+      expect do
+        described_class.raise_if_new_database_connection do
+          ApplicationRecord.connection_pool
+        end
+      end.not_to raise_error
+
       expect do
-        described_class.with_disabled_database_connections do
-          ApplicationRecord.connects_to database: { writing: :main, reading: :main }
+        described_class.raise_if_new_database_connection do
+          Ci::ApplicationRecord.connection_pool
         end
-      end.to raise_error(/Unxpected connection_pools/)
+      end.not_to raise_error
     end
   end
 end