Disallow the use of X.establish_connection

This adds a RuboCop rule that disallows the usage of
`SomeModel.establish_connection`, as this pattern can break database
configurations and slows down the tests.

See https://gitlab.com/gitlab-org/gitlab/-/issues/338653 for more
information.

config/initializers/active_record_lifecycle.rb

@@ -5,7 +5,7 @@
 if defined?(ActiveRecord::Base) && !Gitlab::Runtime.sidekiq?
   Gitlab::Cluster::LifecycleEvents.on_worker_start do
     ActiveSupport.on_load(:active_record) do
-      ActiveRecord::Base.establish_connection
+      ActiveRecord::Base.establish_connection # rubocop: disable Database/EstablishConnection
 
       Gitlab::AppLogger.debug("ActiveRecord connection established")
     end

config/initializers/database_config.rb

@@ -13,6 +13,6 @@
     # The Geo::TrackingBase model does not yet use connects_to. So,
     # this will not properly support geo: from config/databse.yml
     # file yet. This is ACK of the current state and will be fixed.
-    Geo::TrackingBase.establish_connection(Gitlab::Database.geo_db_config_with_default_pool_size)
+    Geo::TrackingBase.establish_connection(Gitlab::Database.geo_db_config_with_default_pool_size) # rubocop: disable Database/EstablishConnection
   end
 end

ee/lib/gitlab/geo/database_tasks.rb

@@ -132,7 +132,7 @@ def load
             ActiveRecord::Tasks::DatabaseTasks.load_schema(ActiveRecord::Base.configurations.configs_for(env_name: 'test').first, :sql, ENV['SCHEMA'])
           ensure
             if should_reconnect
-              ActiveRecord::Base.establish_connection(Gitlab::Geo::DatabaseTasks.db_config)
+              ActiveRecord::Base.establish_connection(Gitlab::Geo::DatabaseTasks.db_config) # rubocop: disable Database/EstablishConnection
             end
           end
         end
@@ -204,7 +204,7 @@ def set_db_env(settings)
         ActiveRecord::Migrator.migrations_paths = ActiveRecord::Tasks::DatabaseTasks.migrations_paths
         Gitlab::Database::SchemaMigrations::Context.default_schema_migrations_path = settings[:schema_migrations_path]
 
-        ActiveRecord::Base.establish_connection(db_config)
+        ActiveRecord::Base.establish_connection(db_config) # rubocop: disable Database/EstablishConnection
       end
 
       class SeedLoader

lib/tasks/gitlab/db.rake

@@ -170,7 +170,7 @@ namespace :gitlab do
       # the `ActiveRecord::Base.connection` might be switched to another one
       # This is due to `if should_reconnect`:
       # https://github.com/rails/rails/blob/a81aeb63a007ede2fe606c50539417dada9030c7/activerecord/lib/active_record/railties/databases.rake#L622
-      ActiveRecord::Base.establish_connection :main
+      ActiveRecord::Base.establish_connection :main # rubocop: disable Database/EstablishConnection
 
       Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
     end

rubocop/cop/database/establish_connection.rb

+# frozen_string_literal: true
+
+module RuboCop
+  module Cop
+    module Database
+      class EstablishConnection < RuboCop::Cop::Cop
+        MSG = "Don't establish new database connections, as this slows down " \
+          'tests and may result in new connections using an incorrect configuration'
+
+        def_node_matcher :establish_connection?, <<~PATTERN
+          (send (const ...) :establish_connection ...)
+        PATTERN
+
+        def on_send(node)
+          add_offense(node, location: :expression) if establish_connection?(node)
+        end
+      end
+    end
+  end
+end

scripts/insert-rspec-profiling-data

@@ -12,7 +12,7 @@ module RspecProfiling
         # This disables the automatic creation of the database and
         # table. In the future, we may want a way to specify the host of
         # the database to connect so that we can call #install.
-        Result.establish_connection(results_url)
+        Result.establish_connection(results_url) # rubocop: disable Database/EstablishConnection
       end
 
       def results_url

spec/lib/gitlab/database/bulk_update_spec.rb

@@ -101,7 +101,7 @@
     before do
       configuration_hash = ActiveRecord::Base.connection_db_config.configuration_hash
 
-      ActiveRecord::Base.establish_connection(
+      ActiveRecord::Base.establish_connection( # rubocop: disable Database/EstablishConnection
         configuration_hash.merge(prepared_statements: prepared_statements)
       )
     end

spec/rubocop/cop/database/establish_connection_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_relative '../../../../rubocop/cop/database/establish_connection'
+
+RSpec.describe RuboCop::Cop::Database::EstablishConnection do
+  subject(:cop) { described_class.new }
+
+  it 'flags the use of ActiveRecord::Base.establish_connection' do
+    expect_offense(<<~CODE)
+      ActiveRecord::Base.establish_connection
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Don't establish new database [...]
+    CODE
+  end
+
+  it 'flags the use of ActiveRecord::Base.establish_connection with arguments' do
+    expect_offense(<<~CODE)
+      ActiveRecord::Base.establish_connection(:foo)
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Don't establish new database [...]
+    CODE
+  end
+
+  it 'flags the use of SomeModel.establish_connection' do
+    expect_offense(<<~CODE)
+      SomeModel.establish_connection
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Don't establish new database [...]
+    CODE
+  end
+end

spec/support/db_cleaner.rb

@@ -67,7 +67,7 @@ def recreate_all_databases!
     # Migrate each database individually
     with_reestablished_active_record_base do
       all_connection_classes.each do |connection_class|
-        ActiveRecord::Base.establish_connection(connection_class.connection_db_config)
+        ActiveRecord::Base.establish_connection(connection_class.connection_db_config) # rubocop: disable Database/EstablishConnection
 
         ActiveRecord::Tasks::DatabaseTasks.migrate
       end

spec/support_specs/database/multiple_databases_spec.rb

@@ -7,13 +7,13 @@
     context 'when doing establish_connection' do
       context 'on ActiveRecord::Base' do
         it 'raises exception' do
-          expect { ActiveRecord::Base.establish_connection(:main) }.to raise_error /Cannot re-establish/
+          expect { ActiveRecord::Base.establish_connection(:main) }.to raise_error /Cannot re-establish/ # rubocop: disable Database/EstablishConnection
         end
 
         context 'when using with_reestablished_active_record_base' do
           it 'does not raise exception' do
             with_reestablished_active_record_base do
-              expect { ActiveRecord::Base.establish_connection(:main) }.not_to raise_error
+              expect { ActiveRecord::Base.establish_connection(:main) }.not_to raise_error # rubocop: disable Database/EstablishConnection
             end
           end
         end
@@ -25,13 +25,13 @@
         end
 
         it 'raises exception' do
-          expect { Ci::ApplicationRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
+          expect { Ci::ApplicationRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/ # rubocop: disable Database/EstablishConnection
         end
 
         context 'when using with_reestablished_active_record_base' do
           it 'does not raise exception' do
             with_reestablished_active_record_base do
-              expect { Ci::ApplicationRecord.establish_connection(:main) }.not_to raise_error
+              expect { Ci::ApplicationRecord.establish_connection(:main) }.not_to raise_error # rubocop: disable Database/EstablishConnection
             end
           end
         end