ActiveRecordで特定のユーザーに関する全データを抽出する

特定のユーザーのデータを全て抽出したい！

サービス運営をしていると、極まれに特定のユーザーに関する全データを抽出したいことがある。
きっと、水平分割や、物理削除してしまったデータの復旧作業をする日がやってくる。

そんな面倒な作業をすることになったあなたのために、今回は特定ユーザーのデータ抽出・インポート方法を紹介する。

ステップ1. 特定のユーザーに紐づく全てのデータを取得する

ActiveRecordを使っていれば、全てのテーブルはモデル層で関連づけられている。
そのため、Userクラスを起点に関連を辿れば、全ての必要なデータを引っ張ってくることができるはずである。

今回は、自前で書いたactive_record_depth_query.rbを使って、関連するレコードを一発で引っ張ってくる。
これを使うと、普段利用しているpreloadのように引数にArray/Hashを渡すことで、外部キーを探索してレコードを取得することができる。

あっという間に、ユーザーに関する全てのレコードを抽出することができた。楽ちん。

depth_query = ActiveRecordDepthQuery.new(user, [:addresses, { cart: :items, orders: [:line_items, :address }]])
depth_query.each do |relation|
  relation #=> 第二引数で指定したレコードのActiveRecord::Relationが順番に入ってくる
end

ステップ2. 抽出したデータを別のDBにコピーする

ステップ1で抽出したデータを、別のDBにコピーする処理を紹介する。
今回はRails6の複数DBの機能と、activerecord-importを利用したので、環境が違う人は適宜書き換えてください。

下記の処理では、古いDB→新しいDB にデータを移すために、
古いDBからレコードを抽出した後、新しいDBに接続を切り替えてからbulk insertしている。

作業としては単純なので、下記のコードで特定のユーザーのデータを抽出・保存することができた。

require 'activerecord-import'

class RecordsFromBackupDatabase
  # インポート対象の関連一覧
  USER_ASSOCIATIONS = [:addresses, { cart: :items, orders: [:line_items, :address] }].freeze

  # 古いDBの接続先
  OLD_DATABASE_HOST = 'old-database.host'

  # 新しいDBの接続先
  NEW_DATABASE_HOST = 'new-database.host'

  # @param user_ids [Array<Integer>] インポート対象のユーザーID
  # 
  # @return [void]
  def perform(user_ids)
    # 古いDB/新しいDBを設定する。共有するプロセスに影響するのでwebサーバーではなくrails consoleで作業すること
    reconnect_database

    ActiveRecord::Base.transaction do
      # 古いDBに接続する
      connected_to_old_database do
        users = User.where(id: user_ids)

        # ユーザーを先に復旧しておく
        import_relation(users)

        users.find_each do |user|
          # ユーザー毎のデータを復旧する
          restore_user(user)
        end
      end
    end
  end

  private

  def reconnect_database
    # 接続中のDBを切断する
    ActiveRecord::Base.connection.disconnect!

    # writing/readingの2系統を定義して、新旧のDBを切り替えられるようにする
    config = ActiveRecord::Base.connection_pool.spec.config
    ActiveRecord::Base.configurations = {
      "#{Rails.env}": config.merge(host: NEW_DATABASE_HOST, replica: false),
      "#{Rails.env}_replica": config.merge(host: OLD_DATABASE_HOST, replica: true)
    }

    ActiveRecord::Base.connects_to(
      database: {
        writing: Rails.env.to_sym,
        reading: :"#{Rails.env}_replica"
      }
    )
  end

  def connected_to_old_database(&block)
    ActiveRecord::Base.connected_to(role: :reading, &block)
  end

  def connected_to_new_database(&block)
    ActiveRecord::Base.connected_to(role: :writing, &block)
  end

  def restore_user(user)
    # ユーザーのassociationを順に探索して、古いDBのレコードを、新しいDBにインポートする
    ActiveRecordDepthQuery.new(user, USER_ASSOCIATIONS).each do |relation|
      relation.find_in_batches do |records|
        klass = relation.klass

        # 本番のDBにレコードをインポートする
        connected_to_new_database do
          klass.import(records, validate: false)
        end
      end
    end
  end
end

まとめ

今回は、preloadライクなインターフェースでレコードを抽出するActiveRecordDepthQueryと、それを利用したデータの復旧方法を紹介した。
願わくば、今回の処理が再び使う日が来ないことを祈るのみだが、もしもそんなオペレーションが必要になった時は、どうぞお使いください。