1 files changed, 46 insertions, 2 deletions

diff --git a/spec/support/database/multiple_databases.rb b/spec/support/database/multiple_databases.rb
index 05f26e57e9c..25c3b6e74ce 100644
--- a/spec/support/database/multiple_databases.rb
+++ b/spec/support/database/multiple_databases.rb
@@ -2,6 +2,15 @@
 
 module Database
   module MultipleDatabases
+    def run_and_cleanup(example)
+      # Each example may call `migrate!`, so we must ensure we are migrated down every time
+      schema_migrate_down!
+
+      example.run
+
+      delete_from_all_tables!(except: deletion_except_tables)
+    end
+
     def skip_if_multiple_databases_not_setup
       skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
     end
@@ -22,6 +31,21 @@ module Database
       model.establish_connection(new_db_config)
     end
 
+    def ensure_schema_and_empty_tables
+      # Ensure all schemas for both databases are migrated back
+      Gitlab::Database.database_base_models.each do |_, base_model|
+        with_reestablished_active_record_base do
+          reconfigure_db_connection(
+            model: ActiveRecord::Base,
+            config_model: base_model
+          )
+
+          schema_migrate_up!
+          delete_from_all_tables!(except: deletion_except_tables)
+        end
+      end
+    end
+
     # The usage of this method switches temporarily used `connection_handler`
     # allowing full manipulation of ActiveRecord::Base connections without
     # having side effects like:
@@ -87,6 +111,16 @@ module Database
 end
 
 RSpec.configure do |config|
+  # Ensure database versions are memoized to prevent query counts from
+  # being affected by version checks. Note that
+  # Gitlab::Database.check_postgres_version_and_print_warning is called
+  # at startup, but that generates its own
+  # `Gitlab::Database::Reflection` so the result is not memoized by
+  # callers of `ApplicationRecord.database.version`, such as
+  # `Gitlab::Database::AsWithMaterialized.materialized_supported?`.
+  # TODO This can be removed once https://gitlab.com/gitlab-org/gitlab/-/issues/325639 is completed.
+  [ApplicationRecord, ::Ci::ApplicationRecord].each { |record| record.database.version }
+
   config.around(:each, :reestablished_active_record_base) do |example|
     with_reestablished_active_record_base(reconnect: example.metadata.fetch(:reconnect, true)) do
       example.run
@@ -99,7 +133,15 @@ RSpec.configure do |config|
     end
   end
 
+  config.append_after(:context, :migration) do
+    break if recreate_databases_and_seed_if_needed
+
+    ensure_schema_and_empty_tables
+  end
+
   config.around(:each, :migration) do |example|
+    self.class.use_transactional_tests = false
+
     migration_schema = example.metadata[:migration]
     migration_schema = :gitlab_main if migration_schema == true
     base_model = Gitlab::Database.schemas_to_base_models.fetch(migration_schema).first
@@ -112,11 +154,13 @@ RSpec.configure do |config|
           config_model: base_model
         )
 
-        example.run
+        run_and_cleanup(example)
       end
     else
-      example.run
+      run_and_cleanup(example)
     end
+
+    self.class.use_transactional_tests = true
   end
 end