Skip to content

Commit 58c27c9

Browse files
author
Alexei Șerșun
committed
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.
1 parent 0b0bfc4 commit 58c27c9

File tree

2 files changed

+25
-17
lines changed

2 files changed

+25
-17
lines changed

lib/database_cleaner/active_record/deletion.rb

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,13 @@ module DatabaseCleaner
55
module ActiveRecord
66
class Deletion < Truncation
77
def clean
8-
connection.disable_referential_integrity do
9-
if pre_count? && connection.respond_to?(:pre_count_tables)
10-
delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
11-
else
12-
delete_tables(connection, tables_to_clean(connection))
8+
with_all_databases do |connection|
9+
connection.disable_referential_integrity do
10+
if pre_count? && connection.respond_to?(:pre_count_tables)
11+
delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
12+
else
13+
delete_tables(connection, tables_to_clean(connection))
14+
end
1315
end
1416
end
1517
end

lib/database_cleaner/active_record/truncation.rb

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,25 +18,31 @@ def initialize(opts={})
1818
end
1919

2020
def clean
21-
connection.disable_referential_integrity do
22-
if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
23-
connection.pre_count_truncate_tables(tables_to_clean(connection))
24-
else
25-
connection.truncate_tables(tables_to_clean(connection))
21+
with_all_databases do |connection|
22+
connection.disable_referential_integrity do
23+
if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
24+
connection.pre_count_truncate_tables(tables_to_clean(connection))
25+
else
26+
connection.truncate_tables(tables_to_clean(connection))
27+
end
2628
end
2729
end
2830
end
2931

3032
private
3133

32-
def connection
33-
@connection ||= ConnectionWrapper.new(
34-
if ::ActiveRecord.version >= Gem::Version.new("7.2")
35-
connection_class.lease_connection
36-
else
37-
connection_class.connection
34+
def with_all_databases
35+
if ::ActiveRecord.version >= Gem::Version.new("6.1")
36+
connection_class.connection_handler.connection_pools.each do |pool|
37+
pool.with_connection do |connection|
38+
connection = ConnectionWrapper.new(connection)
39+
yield connection
40+
end
3841
end
39-
)
42+
else
43+
connection = ConnectionWrapper.new(connection_class.connection)
44+
yield connection
45+
end
4046
end
4147

4248
def tables_to_clean(connection)

0 commit comments

Comments
 (0)