Skip to content

Commit

Permalink
Add support for horizontal sharding for Rails 6.1+
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Alexei Șerșun committed Jun 16, 2024
1 parent 0b0bfc4 commit 58c27c9
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 17 deletions.
12 changes: 7 additions & 5 deletions lib/database_cleaner/active_record/deletion.rb
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,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
Expand Down
30 changes: 18 additions & 12 deletions lib/database_cleaner/active_record/truncation.rb
Original file line number Diff line number Diff line change
Expand Up @@ -18,25 +18,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)
Expand Down

0 comments on commit 58c27c9

Please sign in to comment.