Support Rails 6 multiple database connections.

[masato-hi]

EN)
Fixed to disconnect all connection_handlers in ActiveRecord 6 and later versions.

JA)
ActiveRecord 6では ActiveRecord::Base#clear_all_connections! 及び ActiveRecord::Base#clear_active_connections! は ActiveRecord::Base#default_connection_handler に処理が移譲されるため、
単一のデータベース接続を扱う場合は正常に動作しますが、複数のデータベース接続を扱う場合はdefault_connection_handler以外のconnection_handlerの接続が切れません。
(多くの場合default_connection_handlerはPrimaryデータベース, それ以外のconnection_handlerはReplicaデータベースへ接続されています)

ActiveRecord::Base#connection_handlers はActiveRecord 6で追加されたため、バージョン6以降とそれ以前で処理を分岐し、複数のデータベース接続においてもコネクションを切断出来るよう修正しました。

[kamipo]

connection_handlers will raise NotImplementedError if legacy_connection_handling = false:

https://github.com/rails/rails/blob/90fe5adad06a930c2c5a2b724373cf680c8a5375/activerecord/lib/active_record/core.rb#L227-L234

rails/rails#40637

[masato-hi]

Thank you for your review.
Probably fixed to work with Rails 6.1.

Do you know a better way to get a list of roles?

[kamipo]

You can refer to the original connection management code:

https://github.com/rails/rails/blob/v6.1.0/activerecord/lib/active_record/query_cache.rb#L42-L56

[kamipo]

We can check the condition in definition time rather than runtime, checking it in runtime have a little overhead for the clear_connections method call.

[kamipo]

No raise NoMethodError since Rails 6.1 have ActiveRecord::Base.legacy_connection_handling.

[masato-hi]

@kamipo
Thanks for the review and reference.

Removed the process when legacy_connection_handling is true because the connection to the replica will continue to remain.

[kamipo]

This is correct for legacy_connection_handling = false but is not always correct for legacy_connection_handling = true.

[masato-hi]

@kamipo
I'm sorry my verification wasn't enough.

For example, if you set ActiveRecord::Base.legacy_connection_handling to true and execute the following code will create a connection to primary in ActiveRecord::Base.connection_handler.

ActiveRecord::Base.connection.execute('SELECT 1 FROM DUAL;')

The connection is out of the management of ActiveRecord::Base.connection_handlers, so running the following code will not disconnect the connection.

ActiveRecord::Base.connection_handlers.each_value do |handler|
  handler.connection_pool_list.each(&:disconnect!)
end

I also need to run ActiveRecord::Base.connection_handler.all_connection_pools.each(&:disconnect!). Are there any concerns about that?

[masato-hi]

@kamipo
The problem I'm having trouble with at the moment is something like the sample code below.

The problem is that one of the three cases is wrong or ActiveRecord::Base#connection is not returning the connection_pool in connection_handlers.

(1): Multiple databases should not be configured when legacy_connection_handling is enabled.
(2): You should not call connection outside the connected_to block when legacy_connection_handling is enabled.
(3): ActionRecord::Base#connected_to should not be called (ApplicationRecord # connected_to should be used).

require 'bundler/inline'

gemfile do
  source 'https://rubygems.org'

  gem 'rspec', require: 'rspec/autorun'
  gem 'rails', '6.1.3.2'
  gem 'sqlite3'
end

require 'active_record'

ActiveRecord::Base.legacy_connection_handling = true

# Point (1)
ActiveRecord::Base.configurations = (YAML.load(
<<-YAML
development:
  primary: &development
    adapter: sqlite3
    database: development.sqlite3
    pool: 5
    timeout: 5000
  primary_replica:
    <<: *development
    replica: true
YAML
))

class ApplicationRecord < ActiveRecord::Base
  self.abstract_class = true

  # Point (1)
  connects_to database: {writing: :primary, reading: :primary_replica}
end

RSpec.describe 'activerecord-refresh_connection' do
  before do
    ActiveRecord::Base.establish_connection(:development)
    ActiveRecord::Base.connection_handlers.each_value do |handler|
      handler.connection_pool_list.each do |pool|
        raise 'expect to disconnected!' if pool.connected?
      end
    end
    raise 'expect to disconnected!' if ApplicationRecord.connected?
  end

  after do
    ActiveRecord::Base.connection_handlers.each_value do |handler|
      handler.connection_pool_list.each(&:disconnect!)
    end
    ActiveRecord::Base.clear_all_connections!
  end

  before 'Connect all connections' do
    # Point (2)
    ApplicationRecord.connection.execute('SELECT 1')

    # Point (3)
    ActiveRecord::Base.connected_to(role: :writing) do
      ApplicationRecord.connection.execute('SELECT 1')
    end
    ActiveRecord::Base.connected_to(role: :reading) do
      ApplicationRecord.connection.execute('SELECT 1')
    end
  end

  context 'Check connected to all connections' do
    it do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        expect(handler.connection_pool_list(:writing).count(&:connected?)).to eq(1)
        expect(handler.connection_pool_list(:reading).count(&:connected?)).to eq(1)
      end
      expect(ApplicationRecord.connected?).to be_truthy
    end
  end

  context 'When only call ActiveRecord::ConnectionAdapters::ConnectionPool#disconnect!' do
    before do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        handler.connection_pool_list.each(&:disconnect!)
      end
    end

    it do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        expect(handler.connection_pool_list(:writing).count(&:connected?)).to eq(0)
        expect(handler.connection_pool_list(:reading).count(&:connected?)).to eq(0)
      end
      expect(ApplicationRecord.connected?).to be_truthy
    end
  end

  context 'When only call ActiveRecord::Base#clear_all_connections!' do
    before do
      ActiveRecord::Base.clear_all_connections!
    end

    it do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        expect(handler.connection_pool_list(:writing).count(&:connected?)).to eq(1)
        expect(handler.connection_pool_list(:reading).count(&:connected?)).to eq(1)
      end
      expect(ApplicationRecord.connected?).to be_falsey
    end
  end

  context 'When call all disconnect methods.' do
    before do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        handler.connection_pool_list.each(&:disconnect!)
      end
      ActiveRecord::Base.clear_all_connections!
    end

    it do
      ActiveRecord::Base.connection_handlers.each_value do |handler|
        expect(handler.connection_pool_list(:writing).count(&:connected?)).to eq(0)
        expect(handler.connection_pool_list(:reading).count(&:connected?)).to eq(0)
      end
      expect(ApplicationRecord.connected?).to be_falsey
    end
  end
end