railsでseedデータを簡単に取り込む

本記事の動作環境は以下の通りです。

ruby 2.1.2p95
rails 3.2.17

seedデータの取り込みがカオスになっていた問題を解消した事例を紹介します。

db/seed_datas/に配置したcsvをシードデータとしてデータベースに取り込んでいます。

例えばこんなcsvをdelete-insertしたい場合

id	name
1	ほげ
2	ふが

以前は各テーブルごとに下のような記述をして取り込んでいました。

seeds.rb

require 'csv'

TestTables.delete_all
test_tables_csv = CSV.readlines("db/seed_datas/test_tables.csv")
test_tables_csv.shift # 1行目はフィールド名なのでとばす
test_tables_csv.each do |row|
  record      = TestTables.new
  record.id   = row[0]
  record.name = row[1]
  record.save!
end

テーブルひとつひとつに対して毎回この記述を書くのは面倒ですし、
事故も発生するので、下のようにすっきり記述するためのモジュールを作りました。

seeds.rb

include SeedModule

SeedModule.import("test_tables", delete_all: true)

モジュールの中身はこんな感じ
（実際にはこの2倍くらいコードがあるんですが外部公表用に削っています。）

seed_module.rb

require 'csv'

module SeedModule

  DEFAULT_ENCODING  = 'utf-8'
  DEFAULT_SEED_PATH = './db/seed_datas/'

  @@encoding     = DEFAULT_ENCODING
  @@seed_path    = DEFAULT_SEED_PATH
  @@table_name   = nil
  @@model        = nil

  def self.import(table_name, before_options = {}, after_options = {}, import_options = {})
    @@table_name = table_name
    @@model      = Module.const_get(table_name.camelcase)
    @@encoding   = import_options[:encoding] || DEFAULT_ENCODING

    exec_options(before_options, :BEFORE) if before_options.present?

    import_all

    exec_options(after_options, :AFTER) if after_options.present?
  end

  def self.exec_options(options, timing)
    options.each do |option, value|
      case option
        when :truncate
          next unless value # valueがtrueの場合のみ実行する
          sql = "TRUNCATE TABLE "+@@table_name+";"
          ActiveRecord::Base.connection.execute(sql)
          msg = "全てのレコードを削除:"+sql
        when :delete_all
          next unless value # valueがtrueの場合のみ実行する
          @@model.delete_all
          msg = "全てのレコードを削除"
        when :delete_before
          @@model.delete_all(["id <= " + value.to_s])
          msg = "ID#{value}以下のレコードを削除"
        else
          raise "[#{@@table_name}]#{timing} : 存在しないオプション #{option}"
      end

      puts "[#{@@table_name}]#{timing} : #{msg}"
    end
  end
  private_class_method :exec_options

  def self.import_all
    rows, fields = retrieve_rows_and_fields_from_csv

    rows.each do |row|
      next if row[0].blank? # IDの列が空の行は無視する
      record = @@model.find_or_initialize_by_id(row[0])
      fields.each_with_index do |field, i|
        next if field.blank?
        record.send(field+"=", row[i])
      end
      record.save!
    end

    puts "[#{@@table_name}]Import : #{rows.size} records"
  end
  private_class_method :import_all

  def self.retrieve_rows_and_fields_from_csv
    rows = CSV.read(@@seed_path + @@table_name + ".csv", encoding: @@encoding)
    # CSVの１行目からフィールド名を取得して取り除く
    fields = rows.shift

    return [rows, fields]
  end
  private_class_method :retrieve_rows_and_fields_from_csv

end