Ruby on Rails 6.0.4 における Active Record の複数 DB の仕組み

今携わっている 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 を使う際は

1. ロールに対応するコネクションハンドラーを選択
2. クラス名（または primary）でコネクションプールを選択
3. コネクションプールから 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 は以下のような動作をする。

1. クラスにクラスのインスタンス変数 @connection_specification_name が定義されていればそれを返す
2. 自身が ActiveRecord::Base なら primary を返す
3. 自身が 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 接続でのトラブル調査ではこちらのほうが便利かもしれない。