Fixing Database/MultipleDatabases offenses in specs

We still have many ignored Rubocop Offenses for
Database/MultipleDatabases in specs. This is a MR
to address validate and fix all of them.

Addressing Issue: #362031

Changelog: other

.rubocop_todo/database/multiple_databases.yml

@@ -9,20 +9,3 @@ Database/MultipleDatabases:
     - 'lib/backup/manager.rb'
     - 'lib/gitlab/background_migration/backfill_projects_with_coverage.rb'
     - 'lib/gitlab/background_migration/copy_ci_builds_columns_to_security_scans.rb'
-    - 'spec/db/schema_spec.rb'
-    - 'spec/initializers/database_config_spec.rb'
-    - 'spec/lib/backup/manager_spec.rb'
-    - 'spec/lib/gitlab/database_spec.rb'
-    - 'spec/lib/gitlab/metrics/subscribers/active_record_spec.rb'
-    - 'spec/lib/gitlab/profiler_spec.rb'
-    - 'spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/constraints_spec.rb'
-    - 'spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/joins_spec.rb'
-    - 'spec/support/caching.rb'
-    - 'spec/support/gitlab/usage/metrics_instrumentation_shared_examples.rb'
-    - 'spec/support/helpers/database/database_helpers.rb'
-    - 'spec/support/helpers/database/table_schema_helpers.rb'
-    - 'spec/support/helpers/migrations_helpers.rb'
-    - 'spec/support/helpers/query_recorder.rb'
-    - 'spec/support/helpers/usage_data_helpers.rb'
-    - 'spec/tasks/gitlab/backup_rake_spec.rb'
-    - 'spec/tasks/gitlab/db_rake_spec.rb'

.rubocop_todo/layout/line_length.yml

@@ -5262,8 +5262,6 @@ Layout/LineLength:
     - 'spec/lib/gitlab/usage/metrics/instrumentations/redis_metric_spec.rb'
     - 'spec/lib/gitlab/usage/metrics/name_suggestion_spec.rb'
     - 'spec/lib/gitlab/usage/metrics/names_suggestions/generator_spec.rb'
-    - 'spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/constraints_spec.rb'
-    - 'spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/joins_spec.rb'
     - 'spec/lib/gitlab/usage/metrics/query_spec.rb'
     - 'spec/lib/gitlab/usage/service_ping/payload_keys_processor_spec.rb'
     - 'spec/lib/gitlab/usage/service_ping_report_spec.rb'

spec/db/schema_spec.rb

@@ -6,7 +6,6 @@
 RSpec.describe 'Database schema' do
   prepend_mod_with('DB::SchemaSupport')
 
-  let(:connection) { ActiveRecord::Base.connection }
   let(:tables) { connection.tables }
   let(:columns_name_with_jsonb) { retrieve_columns_name_with_jsonb }
 
@@ -100,50 +99,56 @@
   }.with_indifferent_access.freeze
 
   context 'for table' do
-    ActiveRecord::Base.connection.tables.sort.each do |table|
-      describe table do
-        let(:indexes) { connection.indexes(table) }
-        let(:columns) { connection.columns(table) }
-        let(:foreign_keys) { connection.foreign_keys(table) }
-        let(:loose_foreign_keys) { Gitlab::Database::LooseForeignKeys.definitions.group_by(&:from_table).fetch(table, []) }
-        let(:all_foreign_keys) { foreign_keys + loose_foreign_keys }
-        # take the first column in case we're using a composite primary key
-        let(:primary_key_column) { Array(connection.primary_key(table)).first }
-
-        context 'all foreign keys' do
-          # for index to be effective, the FK constraint has to be at first place
-          it 'are indexed' do
-            first_indexed_column = indexes.map(&:columns).map do |columns|
-              # In cases of complex composite indexes, a string is returned eg:
-              # "lower((extern_uid)::text), group_id"
-              columns = columns.split(',') if columns.is_a?(String)
-              columns.first.chomp
+    Gitlab::Database::EachDatabase.each_database_connection do |connection, _|
+      schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
+
+      connection.tables.each do |table|
+        next unless schemas_for_connection.include?(Gitlab::Database::GitlabSchema.tables_to_schema[table])
+
+        describe table do
+          let(:indexes) { connection.indexes(table) }
+          let(:columns) { connection.columns(table) }
+          let(:foreign_keys) { connection.foreign_keys(table) }
+          let(:loose_foreign_keys) { Gitlab::Database::LooseForeignKeys.definitions.group_by(&:from_table).fetch(table, []) }
+          let(:all_foreign_keys) { foreign_keys + loose_foreign_keys }
+          # take the first column in case we're using a composite primary key
+          let(:primary_key_column) { Array(connection.primary_key(table)).first }
+
+          context 'all foreign keys' do
+            # for index to be effective, the FK constraint has to be at first place
+            it 'are indexed' do
+              first_indexed_column = indexes.map(&:columns).map do |columns|
+                # In cases of complex composite indexes, a string is returned eg:
+                # "lower((extern_uid)::text), group_id"
+                columns = columns.split(',') if columns.is_a?(String)
+                columns.first.chomp
+              end
+              foreign_keys_columns = all_foreign_keys.map(&:column)
+              required_indexed_columns = foreign_keys_columns - ignored_index_columns(table)
+
+              # Add the primary key column to the list of indexed columns because
+              # postgres and mysql both automatically create an index on the primary
+              # key. Also, the rails connection.indexes() method does not return
+              # automatically generated indexes (like the primary key index).
+              first_indexed_column.push(primary_key_column)
+
+              expect(first_indexed_column.uniq).to include(*required_indexed_columns)
             end
-            foreign_keys_columns = all_foreign_keys.map(&:column)
-            required_indexed_columns = foreign_keys_columns - ignored_index_columns(table)
-
-            # Add the primary key column to the list of indexed columns because
-            # postgres and mysql both automatically create an index on the primary
-            # key. Also, the rails connection.indexes() method does not return
-            # automatically generated indexes (like the primary key index).
-            first_indexed_column.push(primary_key_column)
-
-            expect(first_indexed_column.uniq).to include(*required_indexed_columns)
           end
-        end
 
-        context 'columns ending with _id' do
-          let(:column_names) { columns.map(&:name) }
-          let(:column_names_with_id) { column_names.select { |column_name| column_name.ends_with?('_id') } }
-          let(:foreign_keys_columns) { all_foreign_keys.map(&:column).uniq } # we can have FK and loose FK present at the same time
-          let(:ignored_columns) { ignored_fk_columns(table) }
+          context 'columns ending with _id' do
+            let(:column_names) { columns.map(&:name) }
+            let(:column_names_with_id) { column_names.select { |column_name| column_name.ends_with?('_id') } }
+            let(:foreign_keys_columns) { all_foreign_keys.map(&:column).uniq } # we can have FK and loose FK present at the same time
+            let(:ignored_columns) { ignored_fk_columns(table) }
 
-          it 'do have the foreign keys' do
-            expect(column_names_with_id - ignored_columns).to match_array(foreign_keys_columns)
-          end
+            it 'do have the foreign keys' do
+              expect(column_names_with_id - ignored_columns).to match_array(foreign_keys_columns)
+            end
 
-          it 'and having foreign key are not in the ignore list' do
-            expect(ignored_columns).to match_array(ignored_columns - foreign_keys)
+            it 'and having foreign key are not in the ignore list' do
+              expect(ignored_columns).to match_array(ignored_columns - foreign_keys)
+            end
           end
         end
       end
@@ -259,13 +264,13 @@ def get_schemas
 
   context 'primary keys' do
     it 'expects every table to have a primary key defined' do
-      connection = ActiveRecord::Base.connection
+      Gitlab::Database::EachDatabase.each_database_connection do |connection, _|
+        problematic_tables = connection.tables.select do |table|
+          !connection.primary_key(table).present?
+        end.map(&:to_sym)
 
-      problematic_tables = connection.tables.select do |table|
-        !connection.primary_key(table).present?
-      end.map(&:to_sym)
-
-      expect(problematic_tables).to be_empty
+        expect(problematic_tables).to be_empty
+      end
     end
   end
 

spec/initializers/database_config_spec.rb

@@ -7,15 +7,15 @@
     load Rails.root.join('config/initializers/database_config.rb')
   end
 
-  it 'retains the correct database name for the connection' do
-    previous_db_name = ApplicationRecord.connection.pool.db_config.name
-
-    subject
+  it 'retains the correct database name for the main database connection' do
+    expect { subject }.not_to change { ApplicationRecord.connection.pool.db_config.name }
+  end
 
-    expect(ApplicationRecord.connection.pool.db_config.name).to eq(previous_db_name)
+  it 'retains the correct database name for the ci database connection' do
+    expect { subject }.not_to change { Ci::ApplicationRecord.connection.pool.db_config.name }
   end
 
   it 'does not overwrite custom pool settings' do
-    expect { subject }.not_to change { ActiveRecord::Base.connection_db_config.pool }
+    expect { subject }.not_to change { ApplicationRecord.connection_db_config.pool }
   end
 end

spec/lib/backup/manager_spec.rb

@@ -164,7 +164,7 @@
 
     before do
       stub_env('INCREMENTAL', incremental_env)
-      allow(ActiveRecord::Base.connection).to receive(:reconnect!)
+      allow(ApplicationRecord.connection).to receive(:reconnect!)
       allow(Gitlab::BackupLogger).to receive(:info)
       allow(Kernel).to receive(:system).and_return(true)
 

spec/lib/gitlab/database_spec.rb

@@ -154,7 +154,7 @@
   describe '.db_config_for_connection' do
     context 'when the regular connection is used' do
       it 'returns db_config' do
-        connection = ActiveRecord::Base.retrieve_connection
+        connection = ApplicationRecord.retrieve_connection
 
         expect(described_class.db_config_for_connection(connection)).to eq(connection.pool.db_config)
       end
@@ -163,11 +163,11 @@
     context 'when the connection is LoadBalancing::ConnectionProxy', :database_replica do
       it 'returns primary db config even if ambiguous queries default to replica' do
         Gitlab::Database::LoadBalancing::Session.current.use_primary!
-        primary_config = described_class.db_config_for_connection(ActiveRecord::Base.connection)
+        primary_config = described_class.db_config_for_connection(ApplicationRecord.connection)
 
         Gitlab::Database::LoadBalancing::Session.clear_session
         Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-          expect(described_class.db_config_for_connection(ActiveRecord::Base.connection)).to eq(primary_config)
+          expect(described_class.db_config_for_connection(ApplicationRecord.connection)).to eq(primary_config)
         end
       end
     end
@@ -196,9 +196,9 @@
 
     context 'when replicas are configured', :database_replica do
       it 'returns the name for a replica' do
-        replica = ActiveRecord::Base.load_balancer.host
+        main_replica = ApplicationRecord.load_balancer.host
 
-        expect(described_class.db_config_name(replica)).to eq('main_replica')
+        expect(described_class.db_config_name(main_replica)).to eq('main_replica')
       end
     end
   end
@@ -207,7 +207,7 @@
     let(:expected) { %w[foo bar] }
 
     it 'includes only main by default' do
-      allow(::ActiveRecord::Base).to receive(:configurations).and_return(
+      allow(::ActiveRecord::Base).to receive(:configurations).and_return( # rubocop:disable Database/MultipleDatabases
         double(configs_for: %w[foo bar].map { |x| double(name: x) })
       )
 
@@ -215,7 +215,7 @@
     end
 
     it 'excludes geo when that is included' do
-      allow(::ActiveRecord::Base).to receive(:configurations).and_return(
+      allow(::ActiveRecord::Base).to receive(:configurations).and_return( # rubocop:disable Database/MultipleDatabases
         double(configs_for: %w[foo bar geo].map { |x| double(name: x) })
       )
 
@@ -229,13 +229,16 @@
       expect(described_class.gitlab_schemas_for_connection(Ci::Build.connection)).to include(:gitlab_ci, :gitlab_shared)
     end
 
+    # rubocop:disable Database/MultipleDatabases
     it 'does return gitlab_ci when a ActiveRecord::Base is using CI connection' do
       with_reestablished_active_record_base do
         reconfigure_db_connection(model: ActiveRecord::Base, config_model: Ci::Build)
 
-        expect(described_class.gitlab_schemas_for_connection(ActiveRecord::Base.connection)).to include(:gitlab_ci, :gitlab_shared)
+        expect(described_class.gitlab_schemas_for_connection(ActiveRecord::Base.connection))
+          .to include(:gitlab_ci, :gitlab_shared)
       end
     end
+    # rubocop:enable Database/MultipleDatabases
 
     context "when there's CI connection", :request_store do
       before do
@@ -280,7 +283,7 @@
 
     it 'does return empty for non-adopted connections' do
       new_connection = ActiveRecord::Base.postgresql_connection(
-        ActiveRecord::Base.connection_db_config.configuration_hash)
+        ActiveRecord::Base.connection_db_config.configuration_hash) # rubocop:disable Database/MultipleDatabases
 
       expect(described_class.gitlab_schemas_for_connection(new_connection)).to be_nil
     ensure
@@ -404,7 +407,7 @@ def subscribe_events
     context 'within a transaction block' do
       it 'publishes a transaction event' do
         events = subscribe_events do
-          ActiveRecord::Base.transaction do
+          ApplicationRecord.transaction do
             User.first
           end
         end
@@ -423,10 +426,11 @@ def subscribe_events
     context 'within an empty transaction block' do
       it 'publishes a transaction event' do
         events = subscribe_events do
-          ActiveRecord::Base.transaction {}
+          ApplicationRecord.transaction {}
+          Ci::ApplicationRecord.transaction {}
         end
 
-        expect(events.length).to be(1)
+        expect(events.length).to be(2)
 
         event = events.first
         expect(event).not_to be_nil
@@ -440,9 +444,9 @@ def subscribe_events
     context 'within a nested transaction block' do
       it 'publishes multiple transaction events' do
         events = subscribe_events do
-          ActiveRecord::Base.transaction do
-            ActiveRecord::Base.transaction do
-              ActiveRecord::Base.transaction do
+          ApplicationRecord.transaction do
+            ApplicationRecord.transaction do
+              ApplicationRecord.transaction do
                 User.first
               end
             end
@@ -464,7 +468,7 @@ def subscribe_events
     context 'within a cancelled transaction block' do
       it 'publishes multiple transaction events' do
         events = subscribe_events do
-          ActiveRecord::Base.transaction do
+          ApplicationRecord.transaction do
             User.first
             raise ActiveRecord::Rollback
           end

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -7,7 +7,7 @@
 
   let(:env) { {} }
   let(:subscriber) { described_class.new }
-  let(:connection) { ActiveRecord::Base.retrieve_connection }
+  let(:connection) { ApplicationRecord.retrieve_connection }
   let(:db_config_name) { ::Gitlab::Database.db_config_name(connection) }
 
   describe '.load_balancing_metric_counter_keys' do

spec/lib/gitlab/profiler_spec.rb

@@ -144,8 +144,8 @@
       it 'uses the replacement logger for the duration of the block' do
         expect(null_logger).to receive(:debug).and_call_original
 
-        expect { described_class.with_custom_logger(null_logger) { ActiveRecord::Base.logger.debug('foo') } }
-          .to not_change { ActiveRecord::Base.logger }
+        expect { described_class.with_custom_logger(null_logger) { ApplicationRecord.logger.debug('foo') } }
+          .to not_change { ApplicationRecord.logger }
           .and not_change { ActionController::Base.logger }
           .and not_change { ActiveSupport::LogSubscriber.colorize_logging }
       end
@@ -162,7 +162,7 @@
 
       it 'does not modify the standard Rails loggers' do
         expect { described_class.with_custom_logger(nil) {} }
-          .to not_change { ActiveRecord::Base.logger }
+          .to not_change { ApplicationRecord.logger }
           .and not_change { ActionController::Base.logger }
           .and not_change { ActiveSupport::LogSubscriber.colorize_logging }
       end

spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/constraints_spec.rb

@@ -4,11 +4,16 @@
 
 RSpec.describe Gitlab::Usage::Metrics::NamesSuggestions::RelationParsers::Constraints do
   describe '#accept' do
-    let(:collector) { Arel::Collectors::SubstituteBinds.new(ActiveRecord::Base.connection, Arel::Collectors::SQLString.new) }
+    let(:collector) do
+      Arel::Collectors::SubstituteBinds.new(ApplicationRecord.connection, Arel::Collectors::SQLString.new)
+    end
 
     it 'builds correct constraints description' do
       table = Arel::Table.new('records')
-      arel = table.from.project(table['id'].count).where(table[:attribute].eq(true).and(table[:some_value].gt(5)))
+      arel = table.from.project(table['id'].count)
+               .where(table[:attribute]
+               .eq(true)
+               .and(table[:some_value].gt(5)))
       described_class.new(ApplicationRecord.connection).accept(arel, collector)
 
       expect(collector.value).to eql '(records.attribute = true AND records.some_value > 5)'

spec/lib/gitlab/usage/metrics/names_suggestions/relation_parsers/joins_spec.rb

@@ -4,7 +4,9 @@
 
 RSpec.describe Gitlab::Usage::Metrics::NamesSuggestions::RelationParsers::Joins do
   describe '#accept' do
-    let(:collector) { Arel::Collectors::SubstituteBinds.new(ActiveRecord::Base.connection, Arel::Collectors::SQLString.new) }
+    let(:collector) do
+      Arel::Collectors::SubstituteBinds.new(ApplicationRecord.connection, Arel::Collectors::SQLString.new)
+    end
 
     context 'with join added via string' do
       it 'collects join parts' do
@@ -33,7 +35,10 @@
 
         result = described_class.new(ApplicationRecord.connection).accept(arel)
 
-        expect(result).to match_array [{ source: "joins", constraints: "records.id = joins.records_id" }, { source: "second_level_joins", constraints: "joins.id = second_level_joins.joins_id" }]
+        expect(result).to match_array [
+          { source: "joins", constraints: "records.id = joins.records_id" },
+          { source: "second_level_joins", constraints: "joins.id = second_level_joins.joins_id" }
+        ]
       end
     end
   end

spec/support/caching.rb

@@ -37,7 +37,7 @@
   end
 
   config.around(:each, :use_sql_query_cache) do |example|
-    ActiveRecord::Base.cache do
+    ApplicationRecord.cache do
       example.run
     end
   end

spec/support/gitlab/usage/metrics_instrumentation_shared_examples.rb

@@ -30,7 +30,7 @@
   end
 
   before do
-    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+    allow(ApplicationRecord.connection).to receive(:transaction_open?).and_return(false)
   end
 
   it 'has correct generate query' do

spec/support/helpers/database/database_helpers.rb

@@ -5,7 +5,7 @@ module DatabaseHelpers
     # In order to directly work with views using factories,
     # we can swapout the view for a table of identical structure.
     def swapout_view_for_table(view)
-      ActiveRecord::Base.connection.execute(<<~SQL.squish)
+      ApplicationRecord.connection.execute(<<~SQL.squish)
         CREATE TABLE #{view}_copy (LIKE #{view});
         DROP VIEW #{view};
         ALTER TABLE #{view}_copy RENAME TO #{view};
@@ -32,7 +32,7 @@ def with_statement_timeout(timeout)
 
       raise ArgumentError, 'Using a timeout of `0` means to disable statement timeout.' if timeout == 0
 
-      previous_timeout = ActiveRecord::Base.connection
+      previous_timeout = ApplicationRecord.connection
         .exec_query('SHOW statement_timeout')[0].fetch('statement_timeout')
 
       set_statement_timeout("#{timeout}ms")
@@ -61,7 +61,7 @@ def with_statement_timeout(timeout)
     #   set_statement_timeout(0.1)
     #   model.select('pg_sleep(0.11)')
     def set_statement_timeout(timeout)
-      ActiveRecord::Base.connection.execute(
+      ApplicationRecord.connection.execute(
         format(%(SET LOCAL statement_timeout = '%s'), timeout)
       )
     end

spec/support/helpers/database/table_schema_helpers.rb

@@ -3,7 +3,7 @@
 module Database
   module TableSchemaHelpers
     def connection
-      ActiveRecord::Base.connection
+      ActiveRecord::Base.connection # rubocop:disable Database/MultipleDatabases
     end
 
     def expect_table_to_be_replaced(original_table:, replacement_table:, archived_table:)
@@ -17,7 +17,7 @@ def expect_table_to_be_replaced(original_table:, replacement_table:, archived_ta
       expect(table_oid(replacement_table)).to be_nil
     end
 
-    def expect_table_columns_to_match(expected_column_attributes, table_name)
+    def expect_table_columns_to_match(column_attributes, table_name)
       expect(connection.table_exists?(table_name)).to eq(true)
 
       actual_columns = connection.columns(table_name)

spec/support/helpers/migrations_helpers.rb

@@ -102,7 +102,7 @@ def use_fake_application_settings
     # We stub this way because we can't stub on
     # `current_application_settings` due to `method_missing` is
     # depending on current_application_settings...
-    allow(ActiveRecord::Base.connection)
+    allow(ApplicationRecord.connection)
       .to receive(:active?)
       .and_return(false)
     allow(Gitlab::Runtime)
@@ -156,10 +156,10 @@ def disable_migrations_output
   end
 
   def migrate!
-    open_transactions = ActiveRecord::Base.connection.open_transactions
+    open_transactions = ActiveRecord::Base.connection.open_transactions # rubocop:disable Database/MultipleDatabases
     allow_next_instance_of(described_class) do |migration|
       allow(migration).to receive(:transaction_open?) do
-        ActiveRecord::Base.connection.open_transactions > open_transactions
+        ActiveRecord::Base.connection.open_transactions > open_transactions # rubocop:disable Database/MultipleDatabases
       end
     end
 

spec/support/helpers/query_recorder.rb

@@ -19,7 +19,7 @@ def initialize(skip_cached: true, skip_schema_queries: true, log_file: nil, quer
 
     def record(&block)
       # force replacement of bind parameters to give tests the ability to check for ids
-      ActiveRecord::Base.connection.unprepared_statement do
+      ActiveRecord::Base.connection.unprepared_statement do # rubocop:disable Database/MultipleDatabases
         ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
       end
     end

spec/support/helpers/usage_data_helpers.rb

@@ -161,7 +161,7 @@ module UsageDataHelpers
     ).freeze
 
   def stub_usage_data_connections
-    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+    allow(ApplicationRecord.connection).to receive(:transaction_open?).and_return(false)
     allow(::Ci::ApplicationRecord.connection).to receive(:transaction_open?).and_return(false) if ::Ci::ApplicationRecord.connection_class?
 
     allow(Gitlab::Prometheus::Internal).to receive(:prometheus_enabled?).and_return(false)

spec/tasks/gitlab/backup_rake_spec.rb

@@ -151,7 +151,7 @@ def reenable_backup_sub_tasks
   describe 'backup' do
     before do
       # This reconnect makes our project fixture disappear, breaking the restore. Stub it out.
-      allow(ActiveRecord::Base.connection).to receive(:reconnect!)
+      allow(ApplicationRecord.connection).to receive(:reconnect!)
     end
 
     let!(:project) { create(:project, :repository) }

spec/tasks/gitlab/db_rake_spec.rb

@@ -22,7 +22,7 @@
 
   describe 'mark_migration_complete' do
     context 'with a single database' do
-      let(:main_model) { ActiveRecord::Base }
+      let(:main_model) { ApplicationRecord }
 
       before do
         skip_if_multiple_databases_are_setup
@@ -135,7 +135,9 @@
 
       context 'when geo is not configured' do
         before do
+          # rubocop:disable Database/MultipleDatabases
           allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for').and_return([main_config])
+          # rubocop:enable Database/MultipleDatabases
         end
 
         context 'when the schema is already loaded' do
@@ -265,8 +267,10 @@
 
       context 'when geo is not configured' do
         before do
+          # rubocop:disable Database/MultipleDatabases
           allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for')
             .and_return([main_config, ci_config])
+          # rubocop:enable Database/MultipleDatabases
         end
 
         context 'when no database has the schema loaded' do
@@ -367,7 +371,7 @@
 
     with_them do
       it 'outputs changed message for automation after operations happen' do
-        allow(ActiveRecord::Base.connection.schema_migration).to receive(:table_exists?).and_return(schema_migration_table_exists)
+        allow(ApplicationRecord.connection.schema_migration).to receive(:table_exists?).and_return(schema_migration_table_exists)
         allow_any_instance_of(ActiveRecord::MigrationContext).to receive(:needs_migration?).and_return(needs_migrations)
         expect { run_rake_task('gitlab:db:unattended') }. to output(/^#{rake_output}$/).to_stdout
       end
@@ -403,7 +407,9 @@
 
     it 'can be executed multiple times within another rake task' do
       expect_multiple_executions_of_task(test_task_name, clean_rake_task, count: 2) do
+        # rubocop:disable Database/MultipleDatabases
         database_count = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).size
+        # rubocop:enable Database/MultipleDatabases
 
         expect_next_instances_of(Gitlab::Database::SchemaCleaner, database_count) do |cleaner|
           expect(cleaner).to receive(:clean).with(output)
@@ -418,7 +424,7 @@
     let(:schemas) { Gitlab::Database::EXTRA_SCHEMAS }
 
     context 'with a single database' do
-      let(:connection) { ActiveRecord::Base.connection }
+      let(:connection) { ApplicationRecord.connection }
 
       before do
         skip_if_multiple_databases_are_setup
@@ -627,7 +633,9 @@ def expect_objects_to_be_dropped(connection)
     end
 
     it 'defaults to main database' do
+      # rubocop:disable Database/MultipleDatabases
       expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(ActiveRecord::Base.connection).and_call_original
+      # rubocop:enable Database/MultipleDatabases
 
       expect do
         run_rake_task('gitlab:db:enqueue_reindexing_action', "[#{index_name}]")
@@ -828,10 +836,12 @@ def expect_objects_to_be_dropped(connection)
     it 'migrate as nonsuperuser check with default username' do
       allow(Rake::Task['db:drop']).to receive(:invoke)
       allow(Rake::Task['db:create']).to receive(:invoke)
+      # rubocop:disable Database/MultipleDatabases
       allow(ActiveRecord::Base).to receive(:configurations).and_return(configurations)
+      allow(ActiveRecord::Base).to receive(:establish_connection).and_return(connection_pool)
+      # rubocop:enable Database/MultipleDatabases
       allow(configurations).to receive(:configs_for).and_return([configuration])
       allow(configuration).to receive(:configuration_hash).and_return(config_hash)
-      allow(ActiveRecord::Base).to receive(:establish_connection).and_return(connection_pool)
 
       expect(config_hash).to receive(:merge).with({ username: 'gitlab' })
       expect(Gitlab::Database).to receive(:check_for_non_superuser)
@@ -843,10 +853,12 @@ def expect_objects_to_be_dropped(connection)
     it 'migrate as nonsuperuser check with specified username' do
       allow(Rake::Task['db:drop']).to receive(:invoke)
       allow(Rake::Task['db:create']).to receive(:invoke)
+      # rubocop:disable Database/MultipleDatabases
       allow(ActiveRecord::Base).to receive(:configurations).and_return(configurations)
+      allow(ActiveRecord::Base).to receive(:establish_connection).and_return(connection_pool)
+      # rubocop:enable Database/MultipleDatabases
       allow(configurations).to receive(:configs_for).and_return([configuration])
       allow(configuration).to receive(:configuration_hash).and_return(config_hash)
-      allow(ActiveRecord::Base).to receive(:establish_connection).and_return(connection_pool)
 
       expect(config_hash).to receive(:merge).with({ username: 'foo' })
       expect(Gitlab::Database).to receive(:check_for_non_superuser)