production db の内容を staging db に import する雑なスクリプト

開発環境のデータをできるだけ本番に近づける - クックパッド開発者ブログみたいなことをやりたかったが、レプリケーション組むとなると大変なので、雑な mysqldump ベースのデータ移行スクリプトを書いてみた。

仕組み

mysqldump で mysqldump -uroot -T /tmp/ --fields-terminated-by="\t" --fields-optionally-enclosed-by="\"" --lines-terminated-by="\n" --fields-escaped-by="" #{database} #{table} のようにして production db のデータを TSV として吐き出して、LOAD DATA LOCAL INFILE #{tsv_file} INTO TABLE #{table} で TSV ファイルを読み込んで staging db にデータを取り込んでいる。

staging db に対する create_table で AUTO_INCREMENT=600000000 のように offset を指定して、id が 600000000 以上なら staging、未満なら production のもの、ということにして取り込み時に production の古いデータを１回消してから入れ直している。

create_table "table_name", options: "AUTO_INCREMENT=600000000" do |t|
  t.integer "user_id",           null: false
  ...
end

EDIT: 結局 mysqldump ではなく、mysql -B -N -e 'select * from #{table}' とすると tab 区切りで出力を受け取れるので、そちらのほうがスクリプト内で扱いやすいと思い、そちらを使う事にした。

スクリプト

mysql2 gem すら使っておらず、ruby の標準ライブラリと mysql コマンドしか使ってないので bundle install いらずだったりする。mysql コマンドではなく mysql2 gem で全部やれば良かったのでは、と思い始めてたりもする。

script/importdb.rb

# importdb.rb: Import production DB into staging DB as cookpad does
#
# cf. http://techlife.cookpad.com/entry/2014/10/03/110806
#
# This script uses only stdlib, so bundle install is not required
#
# MIT License
require 'yaml'
ROOT = ROOT = File.expand_path('../..', __FILE__)
ENV['RAILS_ENV'] ||= 'development'

def run
  db_info = YAML.load_file(File.join(ROOT, 'config/database.yml'))
  if ENV['RAILS_ENV'] == 'production'
    db_from = db_info['production']
    db_to   = db_info['staging']
  else
    db_from = db_info['development']
    db_to   = db_info['test']
  end

  tables = `#{_mysql(db_from)} -B -N -e 'show tables' | egrep -v 'schema_migrations'`.split("\n")
  tables.each do |table|
    $stdout.puts "Loading #{table} ..."
    tsv_file = TableDumper.new(db_from, table).dump
    TableLoader.new(db_to, table).load(tsv_file)
  end
rescue => e
  $stderr.puts "FAILURE: #{e.class} #{e.message}\n  #{e.backtrace.join("\n  ")}"
  exit 1
end

def _mysql(connect_info)
  cmd = "mysql -u#{connect_info['username']}"
  cmd << " -p#{connect_info['password']}" if connect_info['password']
  cmd << " -h#{connect_info['host']}"     if connect_info['host']
  cmd << " #{connect_info['database']}"
  cmd
end

class TableDumper
  attr_accessor :connect_info, :table
  def initialize(connect_info, table)
    @connect_info = connect_info
    @table        = table
  end

  # Dump Production DB contents to a temp file
  #
  # @return [String] path to the output temp file
  def dump
    column_names = fetch_column_names
    rows = fetch_rows
    rows = filter(column_names, rows)
    tmp_file = write(rows)
  end

  private

  def mysql
    _mysql(connect_info)
  end

  def fetch_column_names
    `#{mysql} -B -e 'select * from #{table} LIMIT 1' | head -1`.chomp.split("\t")
  end

  def fetch_rows
    `#{mysql} -B -N -e "select * from #{table}"`.chomp.split("\n").map {|row| row.split("\t") }
  end

  def filter(column_names, rows)
    # Take care of NULL
    rows.each do |row|
      row.each_with_index do |column, idx|
        row[idx] = '\N' if column == 'NULL'
      end
    end
    
    # Do some filtering if you want
    rows
  end

  def write(rows)
    Tempfile.open("importdb_") do |fp|
      rows.each do |row|
        fp.write(row.join("\t") << "\n")
      end
      fp.path
    end
  end
end

class TableLoader
  attr_accessor :connect_info, :table
  OFFSET = 600000000

  def initialize(connect_info, table)
    @connect_info = connect_info
    @table        = table
  end

  # Load the contents of a file into Staging DB
  #
  # with removing old data
  def load(tsv_file)
    delete_old_data
    load_new_data(tsv_file)
  end

  private

  def mysql
    _mysql(connect_info)
  end

  def delete_old_data
    `#{mysql} -e "delete from #{table} where id < #{OFFSET}"`
  end

  def load_new_data(tsv_file)
    `#{mysql} -e "LOAD DATA LOCAL INFILE '#{tsv_file}' INTO TABLE #{table}"`
  end
end

run

おわりに

とても雑なので参考までに

追記 (2016/01/26)

mysql2 版もつくりました > https://gist.github.com/sonots/d0161b3f4aeaf8cccb48