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](https://github.com/thiagopradi/octopus) や [SwitchPoint](https://github.com/eagletmt/switch_point) のような gem を導入する必要があった。
しかし 6.0 から[複数 DB サポート](https://railsguides.jp/active_record_multiple_databases.html)が入り、別途 gem をインストールすることなく複数の DB にアクセスできるようになった。

```yaml: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
```

```ruby
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` をクラスに記載するという設計なので、コネクション設定はクラス単位となる。
しかし親クラスに設定しておけば、それを子孫クラスで共有できる。

```ruby
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`）。

```ruby
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` の戻り値として、以下のような形で使われている。

```ruby: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)` と同じような動作をするコードを書くと以下のようになる。

```ruby
# 以下のコードはどれも（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` についての説明を省略した。
これについて改めてコードを確認しよう。

```ruby: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` の戻り値を返す

```ruby
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` が記載されたクラス名とコネクションプールの関連付けがなされる。

```ruby
# 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` が評価されるときにそのクラスに定義される。

```ruby
class Book < ApplicationRecord
  @connection_specification_name #=> nil

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

  @connection_specification_name #=> "Book"
end
```

`connects_to` の実装を確認すると

```ruby: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` はない。

```ruby: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 つセットするようになっている。

```ruby: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 つ選びセットする。
こうすることでロールの切り替えができるという仕組みだ。

```ruby
# レコードの削除が目的だが、同時に 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` のロード時に実行されるようになっている。

```ruby: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
```

```ruby: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
```

```ruby: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` を使うはずだ。

```ruby
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` にセットしてブロックを実行するという、先のコードと同じ動作をする。
ただし、こちらは終わったらコネクションハンドラーを元に戻す処理が追加される。

```ruby: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
```

```ruby: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
```

```ruby
# 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` が便利。

```yaml
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
```

```ruby
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 接続でのトラブル調査ではこちらのほうが便利かもしれない。