From b89a6bf040d941c42ebaa4f246553386bc229bc9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexei=20=C8=98er=C8=99un?= <alexei.sersun@saltedge.com>
Date: Sun, 16 Jun 2024 20:21:37 +0300
Subject: [PATCH] Add support for horizontal sharding for Rails 6.1+

Since Rails 6.1, it is possible for a model to connect to multiple
databases. A minimal example of such application is:

```ruby
class ApplicationRecord < ActiveRecord::Base
  connects_to shard: {
    defaul: { writing: :primary_db },
    shard_one: { writing: :secondary_db }
  }
end

class User < ApplicationRecord; end

ApplicationRecord.connected_to(shard: :shard_one, role: :writing) do
  User.create!(...) # creates users in secondary_db DB
end

ApplicationRecord.connection_handler.connection_pools.map { |pool|
pool.db_config.configuration_hash[:database] } # [:primary_db,
:secondary_db]
```

With support for multiple databases for a model, one would have something like this in tests:

```ruby
DatabaseCleaner[:active_record, db: ApplicationRecord]

DatabaseCleaner.start
DatabaseCleaner.clean
```

In `.clean`, however, the bug occurs: it doesn't actually
delete or truncate data from the :secondary_db.

To fix the bug, DatabaseCleaner should iterate through _all_ connection pools the model is connected to.
---
 .../active_record/deletion.rb                 | 12 ++++----
 .../active_record/truncation.rb               | 30 +++++++++++--------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/lib/database_cleaner/active_record/deletion.rb b/lib/database_cleaner/active_record/deletion.rb
index ad28212..78bdd01 100644
--- a/lib/database_cleaner/active_record/deletion.rb
+++ b/lib/database_cleaner/active_record/deletion.rb
@@ -2,11 +2,13 @@ module DatabaseCleaner
   module ActiveRecord
     class Deletion < Truncation
       def clean
-        connection.disable_referential_integrity do
-          if pre_count? && connection.respond_to?(:pre_count_tables)
-            delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
-          else
-            delete_tables(connection, tables_to_clean(connection))
+        with_all_databases do |connection|
+          connection.disable_referential_integrity do
+            if pre_count? && connection.respond_to?(:pre_count_tables)
+              delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
+            else
+              delete_tables(connection, tables_to_clean(connection))
+            end
           end
         end
       end
diff --git a/lib/database_cleaner/active_record/truncation.rb b/lib/database_cleaner/active_record/truncation.rb
index 2e044a0..a3f85e7 100644
--- a/lib/database_cleaner/active_record/truncation.rb
+++ b/lib/database_cleaner/active_record/truncation.rb
@@ -17,25 +17,31 @@ def initialize(opts={})
       end
 
       def clean
-        connection.disable_referential_integrity do
-          if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
-            connection.pre_count_truncate_tables(tables_to_clean(connection))
-          else
-            connection.truncate_tables(tables_to_clean(connection))
+        with_all_databases do |connection|
+          connection.disable_referential_integrity do
+            if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
+              connection.pre_count_truncate_tables(tables_to_clean(connection))
+            else
+              connection.truncate_tables(tables_to_clean(connection))
+            end
           end
         end
       end
 
       private
 
-      def connection
-        @connection ||= ConnectionWrapper.new(
-          if ::ActiveRecord.version >= Gem::Version.new("7.2")
-            connection_class.lease_connection
-          else
-            connection_class.connection
+      def with_all_databases
+        if ::ActiveRecord.version >= Gem::Version.new("6.1")
+          connection_class.connection_handler.connection_pools.each do |pool|
+            pool.with_connection do |connection|
+              connection = ConnectionWrapper.new(connection)
+              yield connection
+            end
           end
-        )
+        else
+          connection = ConnectionWrapper.new(connection_class.connection)
+          yield connection
+        end
       end
 
       def tables_to_clean(connection)