Merge branch 'many-dbs-before-all-adapter' into 'master'

Make `DbCleaner` and `BeforeAllAdapter` to support many connections

See merge request gitlab-org/gitlab!67330

ee/app/models/geo/tracking_base.rb

@@ -10,6 +10,12 @@ module Geo
     SecondaryNotConfigured = Class.new(StandardError)
 
     if ::Gitlab::Geo.geo_database_configured?
+      # Mark current model as a `connection_class`
+      # This is the behavior when executing `connects_to`
+      # which indicates that given class is holding a connection object
+      # and should be treated accordingly.
+      self.connection_class = true
+
       establish_connection Rails.configuration.geo_database
     end
 

ee/spec/support/before_all_adapter.rb

-# frozen_string_literal: true
-
-module EE
-  class BeforeAllAdapter
-    def self.begin_transaction
-      TestProf::BeforeAll::Adapters::ActiveRecord.begin_transaction
-
-      ::Geo::BaseRegistry.connection.begin_transaction
-    end
-
-    def self.rollback_transaction
-      TestProf::BeforeAll::Adapters::ActiveRecord.rollback_transaction
-
-      ::Geo::BaseRegistry.connection.rollback_transaction
-    end
-  end
-end

ee/spec/support/db_cleaner.rb

@@ -8,15 +8,5 @@ module EE
     def deletion_except_tables
       super << 'licenses'
     end
-
-    override :setup_database_cleaner
-    def setup_database_cleaner
-      if ::Gitlab::Geo.geo_database_configured?
-        ::DatabaseCleaner[:active_record, { connection: ::Geo::BaseRegistry }]
-        TestProf::BeforeAll.adapter = ::EE::BeforeAllAdapter
-      end
-
-      super
-    end
   end
 end

spec/support/before_all_adapter.rb

+# frozen_string_literal: true
+
+class BeforeAllAdapter # rubocop:disable Gitlab/NamespacedClass
+  def self.all_connection_pools
+    ::ActiveRecord::Base.connection_handler.all_connection_pools
+  end
+
+  def self.begin_transaction
+    self.all_connection_pools.each do |connection_pool|
+      connection_pool.connection.begin_transaction(joinable: false)
+    end
+  end
+
+  def self.rollback_transaction
+    self.all_connection_pools.each do |connection_pool|
+      if connection_pool.connection.open_transactions.zero?
+        warn "!!! before_all transaction has been already rollbacked and " \
+              "could work incorrectly"
+        next
+      end
+
+      connection_pool.connection.rollback_transaction
+    end
+  end
+end
+
+TestProf::BeforeAll.adapter = ::BeforeAllAdapter

spec/support/db_cleaner.rb

 # frozen_string_literal: true
 
 module DbCleaner
+  def all_connection_classes
+    ::ActiveRecord::Base.connection_handler.connection_pool_names.map(&:constantize)
+  end
+
   def delete_from_all_tables!(except: [])
     except << 'ar_internal_metadata'
 
@@ -12,7 +16,9 @@ module DbCleaner
   end
 
   def setup_database_cleaner
-    DatabaseCleaner[:active_record, { connection: ActiveRecord::Base }]
+    all_connection_classes.each do |connection_class|
+      DatabaseCleaner[:active_record, { connection: connection_class }]
+    end
   end
 end