activerecord-multi-tenant で delete できなかったのでコードリーディングした

この記事は SmartHR Advent Calendar 2019 20日目の記事です。

こんにちは、SmartHR でバックエンドエンジニアをやっている kouryou です。
今回は、SmartHR で利用されている activerecord-multi-tenant という gem で予想外の挙動にハマったため、気になってコードリーディングした話を書きたいと思います。
ちなみにコードリーディングした話を選んだ経緯としては、メタプログラミングRubyという有名な本を読んだことが今年一番勉強になったため、その本で得た知識をアウトプットしようと考えたからです。
これまで gem のコードをじっくり読んだことがなかったので、この機会に読んでみました。

activerecord-multi-tenant とは

そもそも activerecord-multi-tenant を使ったことのない人も多いと思いますので、どういった gem なのかを簡単に説明します。
activerecord-multi-tenant は、1つのサービスの中に複数のテナントが同居するマルチテナントサービスを実現するための gem で、企業やチームで利用する SaaS 型のサービスでよく使われます。

使い方

例として、gem をインストール後、customer をテナントとした customer ごとに database を持つアプリケーションを作成します。
モデル層は以下のようになります。

class Customer < ActiveRecord::Base
  # ...
end

class Site < ActiveRecord::Base
  multi_tenant :customer
  # ...
end

これで customer ごとに別れた database に所属している各 site をいい感じに操作できるようになります。
例えば、あるテナントの最初の site を取得したい場合は、以下のようにテナントを設定すると、通常通り Site.first で取得できます。

pry(main)> # 現在のテナントを設定
pry(main)> MultiTenant.current_tenant = Customer.find('xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx')
pry(main)> # このテナントの最初の site を取得
pry(main)> Site.first

ハマった挙動

上記の例は MultiTenant.current_tenant= で現在のテナントを設定していますが、設定しなかった場合でも、データをきちんと特定できていれば、変更や削除などは通常通りできます。
例えば、site の id が uuid だった場合は、以下のように普通に find して destroy できます。

pry(main)> site = Site.find('yyyyyyyy-yyyy-yyyy-yyyy-yyyyyyyyyyyy')
pry(main)> site.destroy

ところが、これが delete の時は上手くいかなかったのです！
SQLを見てみると、以下のようになっていました。

• destroy の場合の SQL (成功ケース)

DELETE FROM "sites" WHERE "sites"."id" = 'yyyyyyyy-yyyy-yyyy-yyyy-yyyyyyyyyyyy' AND "sites"."customer_id" = 'xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx'

• delete の場合の SQL (失敗ケース)

DELETE FROM "sites" WHERE "sites"."id" = 'yyyyyyyy-yyyy-yyyy-yyyy-yyyyyyyyyyyy' AND "sites"."customer_id" IS NULL

destroy の場合は現在のテナントを設定していなくても、勝手に customer_id を埋めてくれていますが、delete の場合はその補完が効かずに NULL となってしまっています。
ただ、現在のテナントを設定すれば delete も同様に上手くいきました。この現象をまとめると以下のようになります。

削除方法\現在のテナントを設定	設定する	設定しない
destroy	○	○
delete	○	×

この予想外の挙動にハマってしまったため、コードリーディングをしてみました。
読む前に予想がつく人もいるかと思いますが、destory と delete の違いは、コールバックを呼ぶかどうかです。
なので、destroy ではコールバックでよしなにやっていた処理が、delete の場合は呼ばれずにおかしくなってしまっているのだろうと当たりを付けて読んでみました。

コードリーディングしてみた

multi_tenant メソッド

先ほど出てきた魔法の言葉 multi_tenant :customer がマルチテナントの設定をしている箇所なので、まずは multi_tenant メソッドを読んでいきましょう。

lib/activerecord-multi-tenant/model_extensions.rb

def multi_tenant(tenant_name, options = {})
  if to_s.underscore.to_sym == tenant_name
    unless MultiTenant.with_write_only_mode_enabled?
      # This is the tenant model itself. Workaround for https://github.com/citusdata/citus/issues/687
      before_create -> { self.id ||= self.class.connection.select_value("SELECT nextval('" + [self.class.table_name, self.class.primary_key, 'seq'].join('_') + "'::regclass)") }
    end
  else
    # 続く

まず最初のif文ですが、中のコメントにも書いてある通り、これは multi_tenant メソッドを呼び出したのが、実際にマルチテナントになるクラス自身だったケースです。to_s メソッドで自身のクラス名を文字列で出し、underscore.to_sym したものと引数 tenant_name が一致する場合となっています。今回の例で言うと、

class Customer < ActiveRecord::Base
  multi_tenant :customer
end

と Customer クラス内で multi_tenant を呼び出さないと true にならないので、このif文には基本的に入らず、else 以降がメインの処理となります。

続いて else 以降の処理を見ていきます。

lib/activerecord-multi-tenant/model_extensions.rb

def multi_tenant(tenant_name, options = {})
  if to_s.underscore.to_sym == tenant_name
    # 省略
  else
    # 省略
    around_destroy -> (record, block) {
      if MultiTenant.current_tenant_id.nil?
        MultiTenant.with(record.public_send(partition_key)) { block.call }
      else
        block.call
      end
    }
  end
end

お、around_destroy がありましたね！ destroy の際に呼ばれるコールバックなので、おそらくここで問題が起きないような処理をされているのでしょう。
では、around_destroy の中身を見ていきましょう。

MultiTenant.current_tenant_id

最初の if 文で MultiTenant.current_tenant_id が nil かどうかを判定しています。
この MultiTenant.current_tenant_id が何かを見てみましょう。

lib/activerecord-multi-tenant/multi_tenant.rb

require 'request_store'

module MultiTenant
  # 省略

  def self.current_tenant_id
    current_tenant_is_id? ? current_tenant : current_tenant.try(:id)
  end

  def self.current_tenant_is_id?
    current_tenant.is_a?(String) || current_tenant.is_a?(Integer)
  end

  def self.current_tenant
    RequestStore.store[:current_tenant]
  end

  # 省略
end

MultiTenant モジュール内で self.current_tenant_id で定義されています。
current_tenant_is_id? が true ならそのまま current_tenant を返し、 false なら current_tenant から id を取得しようとします。
current_tenant_is_id? メソッドは、 current_tenant が文字列または数値かどうかを判定しています。 MultiTenant 内の current_tenant はオブジェクトと id どちらにも対応しているのでしょう。
では、current_tenant は何をしているのでしょうか？
RequestStore.store[:current_tenant] は、request_store というリクエスト毎にグローバルな変数を使えるようにする gem を利用しています。このスレッド内で定義された current_tenant というグローバル変数を取得しています。
ということは、つまり current_tenant グローバル変数をセットしている箇所がどこかにあるはずです。それが、現在のテナントを設定するために使った MultiTenant.current_tenant= だったのです。(オブジェクトしか対応してないと思い込んでいましたが、id も対応してることを初めて知りました)

lib/activerecord-multi-tenant/multi_tenant.rb

  def self.current_tenant=(tenant)
    RequestStore.store[:current_tenant] = tenant
  end

一通り理解できたので、もう一度 around_destroy の中身を見てみましょう。

lib/activerecord-multi-tenant/model_extensions.rb

    around_destroy -> (record, block) {
      if MultiTenant.current_tenant_id.nil?
        MultiTenant.with(record.public_send(partition_key)) { block.call }
      else
        block.call
      end
    }

最初の if 文では、現在のテナントを設定しているかどうかを判定して、設定済みなら渡された block を呼び出すだけで、設定していない場合は何か特別な処理が入っています。

現在のテナントを設定していない場合の処理

if 文の中身 MultiTenant.with が何者かを見てみましょう。

lib/activerecord-multi-tenant/multi_tenant.rb

  def self.with(tenant, &block)
    return block.call if self.current_tenant == tenant
    old_tenant = self.current_tenant
    begin
      self.current_tenant = tenant
      return block.call
    ensure
      self.current_tenant = old_tenant
    end
  end

現在のテナントを設定している場合は、最初の行で block を呼び出すだけで return しています。今回は設定してないので、この先の処理に進みます。
old_tenant に現在のテナントを突っ込んで、引数で渡されてきた tenant を現在のテナントに設定して block を実行し、最後に現在のテナントを元のテナントに切り戻しています。
つまり、一時的に渡されたテナントに切り替えて block の処理を行うためのメソッドということです。

今回 around_destroy 内では引数 tenant として、record.public_send(partition_key) を渡しています。record は実際に destroy の対象となっている record のことなので、今回のケースでいうと site が該当します。
site に partition_key というメソッドを呼び出そうとしていますが、partition_key とは何でしょうか？

lib/activerecord-multi-tenant/model_extensions.rb

        @partition_key = options[:partition_key] || MultiTenant.partition_key(tenant_name)
        partition_key = @partition_key

multi_tenant メソッド内に定義されていました。
今回 options は渡してないので、MultiTenant.partition_key(tenant_name) が partition_key の返す値になっています。

lib/activerecord-multi-tenant/multi_tenant.rb

  def self.partition_key(tenant_name)
    "#{tenant_name.to_s}_id"
  end

MultiTenant.partition_key はシンプルで、引数 tenant に _id の文字列を付け加えているだけです。なので、今回だと tenant に cutomer を渡しているので、customer_id が返却されます。

まとめ

一通り情報が揃ったのでまとめると、現在のテナントを設定せずに destory を呼び出した場合は、
MultiTenant.with(record.public_send(partition_key)) { block.call }
の処理に入ります。ここでは、site.customer_id を呼び出し、その customer_id で一時的に現在のテナントを切り替えて destroy が行われるということです。
最初に予想した通り、現在のテナントを設定していない場合は、コールバックを使ってテナントを設定し直しているため、destroy では削除が可能になっているのでした！
delete ではこのコールバックが呼ばれないので、現在のテナントが設定されず、NULL が入ってしまい、削除できなかったという結論でした。

修正されてた

そして、この問題を見つけてからすぐに、修正が施されていました！(ちなみに私は何もしてません。偶然のタイミングで修正が入っていました)
ということで、最新版ではどのような修正が入ったのか見てみましょう。

lib/activerecord-multi-tenant/persistence_extension.rb

module ActiveRecord
  module Persistence
    alias :delete_orig :delete

    def delete
      if MultiTenant.multi_tenant_model_for_table(self.class.table_name).present? && persisted? && MultiTenant.current_tenant_id.nil?
        MultiTenant.with(self.public_send(self.class.partition_key)) { delete_orig }
      else
        delete_orig
      end
    end
  end
end

オープンクラスとアラウンドエイリアスという手法を使って、元の delete を上書きして使っていました。
後はほぼ around_destroy 内の処理と同じですが、if 文の条件で
MultiTenant.multi_tenant_model_for_table(self.class.table_name).present? && persisted?
が追加で増えています。
persisted? は DB に保存済みかどうか( new しただけで save してないデータとか)を判定しています。
MultiTenant.multi_tenant_model_for_table の中身を見てみましょう。

lib/activerecord-multi-tenant/multi_tenant.rb

  def self.multi_tenant_model_for_table(table_name)
    @@multi_tenant_models ||= {}
    @@multi_tenant_models[table_name.to_s]
  end

@@multi_tenant_models というクラス変数に登録されているテーブル名をキーとして、その値を取得しています。@@multi_tenant_models は

lib/activerecord-multi-tenant/multi_tenant.rb

  def self.register_multi_tenant_model(table_name, model_klass)
    @@multi_tenant_models ||= {}
    @@multi_tenant_models[table_name.to_s] = model_klass
  end

にて挿入されていて、実は multi_tenant メソッド内でこの register_multi_tenant_model は呼び出されていました。

lib/activerecord-multi-tenant/model_extensions.rb

        MultiTenant.register_multi_tenant_model(table_name, self) if table_name

なぜ MultiTenant.multi_tenant_model_for_table(self.class.table_name).present? の判定が必要なのか最初わからなかったのですが、どうやら全テナント共通のテーブルのデータの削除ができない問題があったので、それを回避するためにこの判定が入ったみたいです。
https://github.com/citusdata/activerecord-multi-tenant/commit/c5585df0076533a8fdf754cb1ed4cc332b5cff3b
たしかにオープンクラスなので、元の delete を呼び出したいときにも影響してしまうため、判定が増えそうですね。

おわりに

メタプログラミングRuby を読んで、gem を読み解く力がついたので、来年は色んな gem のコードリーディングをして、Ruby 力を高めていこうと思います