今携わっている Rails 6.0 のプロジェクトで複数 DB に移行する際に、すんなりいかず実装を深く追うことになった。
その際に得た Active Record のコネクションハンドラーまわりについてのメモを残そうと思う。
なお、このエントリは Ruby on Rails 6.0.4 について調べたものであり、6.1 の新機能であるシャーディングなどについては対象外なので注意。

Ruby on Rails 6.0 における複数 DB の概要

従来、Ruby on Rails は 5.2 まで複数 DB への接続をサポートしておらず、Octopus や SwitchPoint のような gem を導入する必要があった。
しかし 6.0 から複数 DB サポートが入り、別途 gem をインストールすることなく複数の DB にアクセスできるようになった。

config/database.yml

default: &default
  adapter: mysql2
  username: root

development:
  main:
    <<: *default
    database: my_main_development

  main_replica:
    <<: *default
    database: my_main_development
    replica: true

  sub:
    <<: *default
    database: my_sub_development
    migrations_paths: db/sub

  sub_replica:
    <<: *default
    database: my_sub_development
    replica: true

class ApplicationRecord < ActiveRecord::Base
  self.abstract_class = true
  connects_to database: { writing: :main, reading: :main_replica }
end

class Author < ApplicationRecord
end

class Book < ApplicationRecord
  connects_to database: { writing: :sub, reading: :sub_replica }
end

connects_to をクラスに記載するという設計なので、コネクション設定はクラス単位となる。
しかし親クラスに設定しておけば、それを子孫クラスで共有できる。

Author.find(1) #=> my_main_development の authors にアクセス
Book.find(1) #=> my_sub_development の books にアクセス

ActiveRecord::Base.connected_to(role: :reading) { Author.find(1) } #=> main_replica の設定で authors にアクセス

複数 DB へのアクセスのしくみ

ActiveRecord::Base.connection とコネクションハンドラー

Author.find(1) は、実装をたどっていくと ActiveRecord::Base.connection.execute(sql) と同等のコードを実行している（ActiveRecord::ConnectionAdapters::MySQL::DatabaseStatements#execute）。

Author.find(1)
# ActiveRecord::Base.connection を経由して SQL を実行する
Author.connection.execute("SELECT `authors`.* FROM `authors` WHERE `authors`.`id` = 1 LIMIT 1")

ActiveRecord::Base.connection は DB と通信をする際に用いる Mysql2::Client を含んだ、ActiveRecord::ConnectionAdapters::Mysql2Adapter のインスタンスを返す（コネクションアダプターが mysql2 の場合）。
複数の DB に対応するために、Active Record はそれに応じた Mysql2Adapter を複数用意し接続先に応じて使い分ける。
この複数のコネクションアダプターの管理を引き受けているのがコネクションハンドラー ActiveRecord::ConnectionAdapters::ConnectionHandler だ。
ActiveRecord::Base.connection_handler の戻り値として、以下のような形で使われている。

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    def connection
      retrieve_connection
    end

    def retrieve_connection
      # この connection_handler は ActiveRecord::Base.connection_handler であり、
      # ActiveRecord::ConnectionAdapters::ConnectionHandler のインスタンスを返す
      connection_handler.retrieve_connection(connection_specification_name)
    end

    # ...
  end
end

コネクションハンドラーの retrieve_connection メソッドを呼んでいる。
これはコネクションハンドラーから ActiveRecord の子孫クラスに対応するコネクションアダプターを得るコードだ。

ではコネクションハンドラーのデータ構造はどのような構造になっているのだろうか。

コネクションハンドラーのデータ構造

コネクションハンドラー ActiveRecord::ConnectionAdapters::ConnectionHandler の、コネクションアダプターについてのデータ構造は以下の通り。

@owner_to_pool
- { プロセスID => Concurrent::Map のインスタンス } の形式の、Concurrent::Map のインスタンス
  - { connection_specification_name => コネクションプール } の形式の、Concurrent::Map のインスタンス
  - "primary" => main（writing ロールの場合。reading ロールは main_replica）向けの ActiveRecord::ConnectionAdapters::ConnectionPool のインスタンス
    - ActiveRecord::ConnectionAdapters::Mysql2Adapter のインスタンス
  - "Book" => sub（writing ロールの場合。reading ロールは sub_replica）向けの ActiveRecord::ConnectionAdapters::ConnectionPool のインスタンス
    - ActiveRecord::ConnectionAdapters::Mysql2Adapter のインスタンス

connects_to を記載したクラス名ごとにコネクションプールを用意している。
ただし、クラスが ApplicationRecord の場合は特別に primary となる。
Author クラスは connects_to を記載していないが、ApplicationRecord を継承しているので primary を使うことになる。

ところで Active Record では ActiveRecord::Base.connected_to を使うことによって writing ロール、reading ロールの切り替えができる。これに対応するデータ構造はどのようになっているかというと

ActiveRecord::Base.connection_handlers
- { ロール => コネクションハンドラー } の形式の Hash
- :writing => writing ロール用の ActiveRecord::ConnectionAdapters::ConnectionHandler のインスタンス
- :reading => reading ロール用の ActiveRecord::ConnectionAdapters::ConnectionHandler のインスタンス

ロールそれぞれにコネクションハンドラーが用意されている。
したがって、Mysql2Adapter を使う際は

ロールに対応するコネクションハンドラーを選択
クラス名（または primary）でコネクションプールを選択
コネクションプールから Mysql2Adapter を取得

という手順を踏む。
これに基づき Author.find(1) と同じような動作をするコードを書くと以下のようになる。

# 以下のコードはどれも（SQL の実行について）同じ
Author.find(1)

Author.connection.
       execute("SELECT `authors`.* FROM `authors` WHERE `authors`.`id` = 1 LIMIT 1")

ActiveRecord::Base.connection_handlers.
                   fetch(:writing).                        # writing ロールのコネクションハンドラー
                   instance_variable_get(:@owner_to_pool). # インスタンス変数 @owner_to_pool
                   fetch(Process.pid).
                   fetch("primary").                       # Author クラスは ApplicationRecord と同じ "primary"
                   connection.                             # ここで ActiveRecord::ConnectionAdapters::Mysql2Adapter が得られる
                   execute("SELECT `authors`.* FROM `authors` WHERE `authors`.`id` = 1 LIMIT 1")

DB の切り替え

先に引用したコードではコネクションハンドラーの retrieve_connection メソッドに渡す connection_specification_name についての説明を省略した。
これについて改めてコードを確認しよう。

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    def connection
      retrieve_connection
    end

    def retrieve_connection
      connection_handler.retrieve_connection(connection_specification_name)
    end

    # 以下 connection_specification_name についての実装
    attr_writer :connection_specification_name

    def connection_specification_name
      if !defined?(@connection_specification_name) || @connection_specification_name.nil?
        return self == Base ? "primary" : superclass.connection_specification_name
      end
      @connection_specification_name
    end
  end
end

先に書いたとおり、connection_handler.retrieve_connection は引数の文字列に対応するコネクションプールを返す。

そして connection_specification_name は以下のような動作をする。

クラスにクラスのインスタンス変数 @connection_specification_name が定義されていればそれを返す
自身が ActiveRecord::Base なら primary を返す
自身が ActiveRecord::Base でなければ親クラスの connection_specification_name の戻り値を返す

ApplicationRecord.connection_specification_name #=> "primary"

# @connection_specification_name が定義されていないので、親クラス（ApplicationRecord）のものが参照される
Author.instance_variable_defined?(:@connection_specification_name) #=> false
Author.connection_specification_name #=> "primary"

# connects_to を記述した Book クラスの名前が @connection_specification_name に入っているので、それを返す
Book.instance_variable_get(:@connection_specification_name) #=> "Book"
Book.connection_specification_name #=> "Book"

親クラスのコネクション設定が子孫クラスにも継承される機能は 3. の動作による。
これによって connects_to が記載されたクラス名とコネクションプールの関連付けがなされる。

# Author.connection が返す Mysql2Adapter は main (my_main_development) 向け
Author.connection.instance_variable_get(:@config)
#=> { ..., database: "my_main_development", ... }

# Book.connection が返す Mysql2Adapter は sub (my_sub_development) 向け
Book.connection.instance_variable_get(:@config)
#=> { ..., database: "my_sub_development", ... }

[補足] `@connection_specification_name` はどこで設定されるのか

結論を先に書いてしまうと、@connection_specification_name は connects_to が評価されるときにそのクラスに定義される。

class Book < ApplicationRecord
  @connection_specification_name #=> nil

  connects_to database: { writing: :sub, reading: :sub_replica }

  @connection_specification_name #=> "Book"
end

connects_to の実装を確認すると

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    def connects_to(database: {})
      connections = []

      database.each do |role, database_key|
        # role には :writing, :reading などが、database_key には :main, :main_replica などが入る
        config_hash = resolve_config_for_connection(database_key)
        handler = lookup_connection_handler(role.to_sym)

        connections << handler.establish_connection(config_hash)
      end

      connections
    end

    def lookup_connection_handler(handler_key)
      handler_key ||= ActiveRecord::Base.writing_role # デフォルトは :writing
      connection_handlers[handler_key] ||= ActiveRecord::ConnectionAdapters::ConnectionHandler.new
    end

    def resolve_config_for_connection(config_or_env)
      raise "Anonymous class is not allowed." unless name

      # config_or_env には :development や :production などの DB 名（複数 DB 以前の使い方）か、
      # :main, :main_replica などのロール名が渡されてくる
      config_or_env ||= DEFAULT_ENV.call.to_sym
      pool_name = primary_class? ? "primary" : name
      # ここで @connection_specification_name がセットされる
      self.connection_specification_name = pool_name

      # config_or_env とクラス名（または "primary"）から DB 設定を特定し、name を追加して返す
      resolver = ConnectionAdapters::ConnectionSpecification::Resolver.new(Base.configurations)
      config_hash = resolver.resolve(config_or_env, pool_name).symbolize_keys
      config_hash[:name] = pool_name

      config_hash
    end

    def primary_class?
      self == Base || defined?(ApplicationRecord) && self == ApplicationRecord
    end
  end
end

connects_to は指定された名前で DB 設定を検索し、その設定に対応するコネクションプールを作成してコネクションハンドラーに登録する。
メソッド名が getter 的なので意表をつかれるが、実はこのうちの DB 設定の検索（resolve_config_for_connection）の際に、setter で @connection_specification_name が設定されるようになっている。

ロールの切り替え

`ActiveRecord::Base.connection_handler` の戻り値

コネクションハンドラーは ActiveRecord::Base.connection_handlers で得られる。
また、DB との通信は ActiveRecord::Base.connection を介して ActiveRecord::Base.connection_handler を通じて行われる。

しかし先のコードには connection_handler はあるが connection_handlers はない。

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    # ...

    def retrieve_connection
      # connection_handler は ActiveRecord::Base.connection_handler
      connection_handler.retrieve_connection(connection_specification_name)
    end

    # ...
  end
end

では ActiveRecord::Base.connection_handlers と ActiveRecord::Base.connection_handler はどのような関係にあるのか。
実は ActiveRecord::Base.connection_handler はスレッドローカル変数へのアクセッサであり、ここにコネクションハンドラーを 1 つセットするようになっている。

activerecord/lib/active_record/core.rb

module ActiveRecord
  module Core
    extend ActiveSupport::Concern

    included do
      class_attribute :default_connection_handler, instance_writer: false

      # getter
      def self.connection_handler
        Thread.current.thread_variable_get("ar_connection_handler") || default_connection_handler
      end

      # setter
      def self.connection_handler=(handler)
        Thread.current.thread_variable_set("ar_connection_handler", handler)
      end

      self.default_connection_handler = ConnectionAdapters::ConnectionHandler.new
    end
  end
end

ここに、ロールに応じて connection_handlers 中のコネクションハンドラーを 1 つ選びセットする。
こうすることでロールの切り替えができるという仕組みだ。

# レコードの削除が目的だが、同時に Author.connects_to を評価することで
# reading ロールのコネクションハンドラー作成もする
Author.delete_all
Author.count #=> 0

# connection_handler に writing ロールのコネクションハンドラーをセットすると、レコードを作成する
sql = "INSERT INTO `authors` (`created_at`, `updated_at`) VALUES (NOW(), NOW())"
ActiveRecord::Base.connection_handler = ActiveRecord::Base.connection_handlers[:writing]
ActiveRecord::Base.connection.execute(sql)
Author.count #=> 1

# connection_handler に reading ロールのコネクションハンドラーをセットすると、レコードは作成できない
ActiveRecord::Base.connection_handler = ActiveRecord::Base.connection_handlers[:reading]
ActiveRecord::Base.connection.execute(sql) #=> ActiveRecord::ReadOnlyError
Author.count #=> 1

[補足] `ActiveRecord::Base.default_connection_handler` について

ar_connection_handler が nil の場合は default_connection_handler が評価される。
ここを見る限り引数なしで単に ConnectionAdapters::ConnectionHandler のインスタンスを作成しているだけで DB の設定がないので、接続ができないように見える。

実はこの設定をするコードは別の場所にあり、on_load フックによって ActiveRecord::Base のロード時に実行されるようになっている。

activerecord/lib/active_record/railtie.rb

module ActiveRecord
  class Railtie < Rails::Railtie
    initializer "active_record.initialize_database" do
      ActiveSupport.on_load(:active_record) do
        self.connection_handlers = { writing_role => ActiveRecord::Base.default_connection_handler }
        # config/database.yml の設定を ActiveRecord にセットする
        self.configurations = Rails.application.config.database_configuration
        establish_connection
      end
    end
  end
end

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    def establish_connection(config_or_env = nil)
      # :development や :test などの環境名か :main、:main_replica、:sub などの設定名に応じた DB 接続設定を得る
      # ActiveRecord::Base のロード時には環境(:development)中の一番最初である :main の設定を返す
      config_hash = resolve_config_for_connection(config_or_env)
      connection_handler.establish_connection(config_hash)
    end
  end
end

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

module ActiveRecord
  module ConnectionAdapters
    class ConnectionHandler
      # メソッド名が ActiveRecord::Base.establish_connection と同名なのでややこしいが、
      # こちらは ActiveRecord::Base.connection_handler.establish_connection の実装
      def establish_connection(config)
        resolver = ConnectionSpecification::Resolver.new(Base.configurations)
        # spec には name が "primary"、config が :main の設定の ActiveRecord::ConnectionAdapters::ConnectionSpecification が代入される
        spec = resolver.spec(config)

        # 〜中略〜 （既存のコネクションの切断処理と、Instrumentation 向けの処理が続く）

        message_bus.instrument("!connection.active_record", payload) do
          # ここでコネクションプールを作成し、"primary" のコネクションプールを @owner_to_pool[Process.pid] にセットする
          owner_to_pool[spec.name] = ConnectionAdapters::ConnectionPool.new(spec)
        end

        owner_to_pool[spec.name]
      end
    end
  end
end

まず writing ロール用のコネクションハンドラーに ActiveRecord::Base.default_connection_handler をセット。
次に config/database.yml の設定を ActiveRecord::Base にロードして ActiveRecord::Base.establish_connection を呼ぶ。
ActiveRecord::Base.connection_handler は何も設定されていない状態では ActiveRecord::Base.default_connection_handler を返すので、これに対して ActiveRecord::ConnectionAdapters::ConnectionHandler#establish_connection が呼ばれ、コネクションプールが作成される。

`connected_to` によるロールの切り替え

ActiveRecord::Base.connection_handler の内容を変更することでロールを切り替えられるのは前述のとおり。
しかし通常は ActiveRecord::Base.connected_to を使うはずだ。

ActiveRecord::Base.connected_to(role: :reading) { Author.create }
#=> ActiveRecord::ReadOnlyError

ActiveRecord::Base.connected_to(role: :writing) { Author.create }
#=> #<Author ...>

ActiveRecord::Base.connected_to は、指定されたロールのコネクションハンドラーを ActiveRecord::Base.connection_handler にセットしてブロックを実行するという、先のコードと同じ動作をする。
ただし、こちらは終わったらコネクションハンドラーを元に戻す処理が追加される。

activerecord/lib/active_record/connection_handling.rb

module ActiveRecord
  module ConnectionHandling
    def connected_to(database: nil, role: nil, prevent_writes: false, &blk)
      if database && role
        # ...
      elsif database
        # ...
      elsif role
        # reading ロールのときは prevent_writes に true をセットする
        # reading 系として登録されていないクエリを実行するとき、prevent_writes が true
        # だと ActiveRecord::ReadOnlyError が上がる
        prevent_writes = true if role == reading_role # デフォルトでは :reading

        with_handler(role.to_sym) do
          connection_handler.while_preventing_writes(prevent_writes, &blk)
        end
      else
        # ...
      end
    end

    def with_handler(handler_key, &blk)
      handler = lookup_connection_handler(handler_key)
      swap_connection_handler(handler, &blk)
    end

    def lookup_connection_handler(handler_key)
      handler_key ||= ActiveRecord::Base.writing_role # デフォルトでは :writing
      connection_handlers[handler_key] ||= ActiveRecord::ConnectionAdapters::ConnectionHandler.new
    end

    private
      # 指定されたコネクションハンドラーと ActiveRecord::Base.connection_handler を swap し、
      # block の評価後にもとに戻す
      def swap_connection_handler(handler, &blk)
        old_handler, ActiveRecord::Base.connection_handler = ActiveRecord::Base.connection_handler, handler
        yield
      ensure
        ActiveRecord::Base.connection_handler = old_handler
      end
  end
end

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

module ActiveRecord
  module ConnectionAdapters
    class ConnectionHandler
      # prevent_writes も connection_handler と同じようにスレッドローカル変数で管理されており、
      # getter と setter がある
      def while_preventing_writes(enabled = true)
        original, self.prevent_writes = self.prevent_writes, enabled
        yield
      ensure
        self.prevent_writes = original
      end
    end
  end
end

# connection_handler には writing ロール用のコネクションハンドラーがセットされている
ActiveRecord::Base.connection_handlers[:writing].object_id #=> 10000
ActiveRecord::Base.connection_handler.object_id #=> 10000 

# 最初は reading ロールのコネクションハンドラーはない
ActiveRecord::Base.connection_handlers[:reading] #=> nil

# writing ロールに切り替えると、コネクションハンドラーが作成・登録され、connection_handler にセットされる
ActiveRecord::Base.connected_to(role: :reading) { ActiveRecord::Base.connection_handler.object_id } #=> 10120
ActiveRecord::Base.connection_handlers[:reading].object_id #=> 10120

# connected_to のブロックを抜けると connection_handler は元に戻る
ActiveRecord::Base.connection_handler.object_id #=> 10000

まとめ

複数の DB に対応するために、Active Record はコネクションハンドラー ActiveRecord::ConnectionAdapters::ConnectionHandler で複数のコネクションアダプターを管理している。
コネクションハンドラーは reading、writing ロールごとに用意されており、connects_to を記載したクラス名ごとのコネクションプールを保持している。
DB への処理は ActiveRecord::Base.connection_handler 経由。connected_to のブロック実行時に、ActiveRecord::Base.connection_handlers 中のコネクションハンドラーのどれか 1 つに差し替えられる。

おまけ

あるコンテキストで使用される DB の設定を確認するには、Model.connection_config が便利。

development:
  main:
    <<: *default
    database: my_main_development
    # connect_timeout で区別をつけられるよう、それぞれの設定に異なる値をセットする
    connect_timeout: 120

  main_replica:
    <<: *default
    database: my_main_development
    connect_timeout: 121
    replica: true

  sub:
    <<: *default
    database: my_sub_development
    migrations_paths: db/sub
    connect_timeout: 122

  sub_replica:
    <<: *default
    database: my_sub_development
    connect_timeout: 123
    replica: true

Author.connection_config
#=> { ..., :database=>"my_main_development", :connect_timeout=>120 }

Book.connection_config
#=> { ..., :database=>"my_sub_development", ...,  :connect_timeout=>122 }

# reading role を指定すると replica の設定が使われる
ActiveRecord::Base.connected_to(role: :reading) { Author.connection_config }
  #=> { ..., :database=>"my_main_development", :connect_timeout=>121, :replica=>true }

Model.connection.pool.spec.config でもほぼ同じ結果を得られる（異常系の処理が多少異なるが、正常系のロジックは同じ）。
DB 接続でのトラブル調査ではこちらのほうが便利かもしれない。