Skip to content

Commit b89a6bf

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 d34716a commit b89a6bf

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
@@ -2,11 +2,13 @@ module DatabaseCleaner
22
module ActiveRecord
33
class Deletion < Truncation
44
def clean
5-
connection.disable_referential_integrity do
6-
if pre_count? && connection.respond_to?(:pre_count_tables)
7-
delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
8-
else
9-
delete_tables(connection, tables_to_clean(connection))
5+
with_all_databases do |connection|
6+
connection.disable_referential_integrity do
7+
if pre_count? && connection.respond_to?(:pre_count_tables)
8+
delete_tables(connection, connection.pre_count_tables(tables_to_clean(connection)))
9+
else
10+
delete_tables(connection, tables_to_clean(connection))
11+
end
1012
end
1113
end
1214
end

lib/database_cleaner/active_record/truncation.rb

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

1919
def clean
20-
connection.disable_referential_integrity do
21-
if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
22-
connection.pre_count_truncate_tables(tables_to_clean(connection))
23-
else
24-
connection.truncate_tables(tables_to_clean(connection))
20+
with_all_databases do |connection|
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))
26+
end
2527
end
2628
end
2729
end
2830

2931
private
3032

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

4147
def tables_to_clean(connection)

0 commit comments

Comments
 (0)