Ruby：使用ActiveRecord验证将CSV数据导入DB

Posted 2021-02-25

This script is a nice balance between validating data and speed.

require 'rubygems'
require 'activerecord'
require 'faster_csv'
 
###########################################
#1. Create a DB table
###########################################
class CreateLocations < ActiveRecord::Migration
  def self.up
    create_table :locations, :force => true do |t|
      t.string :sale_type,:home_type,:address,:city,:state,:zip,:list_price,:beds,:baths,:location,:sqft,:lot_size,:year_built,:parking_spots,:parking_type,:days_on_market,:next_open_house_date,:next_open_house_start_time,:next_open_house_end_time,:recent_reduction_date,:original_list_price,:last_sale_date,:last_sale_price,:url,:source,:listing_id,:original_source,:favorite,:interested,:latitude,:longitude,:is_short_sale
    end
  end
 
  def self.down
    drop_table :locations
  end
end
 
class Location < ActiveRecord::Base
  validates_presence_of #:sale_type,:home_type,:address,:city,:state,:zip,:list_price,:beds,:baths,:location,:sqft,:year_built,:parking_spots,:days_on_market,:original_list_price,:url,:source,:listing_id,:original_source,:favorite,:interested,:is_short_sale
end
 
###########################################
#2. Connect to Database: ar-extensions is the best of both worlds between ActiveRecord and DB default import features
###########################################
require 'ar-extensions'
#Logger Objects are handy for finding out why imports crash
logger = Logger.new('import.log')
#Set the Logger level to Info to prevent boring debug messages
logger.level = Logger::INFO
#Make DB connection
ActiveRecord::Base.establish_connection(
  :adapter  => 'mysql',
  :database => 'test_test',
  :username => 'root',
  :password => ''
)
#Connect Logger to Active Record
ActiveRecord::Base.logger = logger
 
#Kick off DB table creation
CreateLocations.up
 
###########################################
#3. Parse the CSV File into a Location Object
###########################################
def load(filename, chunk_size = 3_000, validate = true)
  options = {
    :headers => true,
    :header_converters => :symbol,
    :col_sep => ','
  }
 
  locations = []
  #Instead of reading all the files at once, we instead break them into chunks
  FasterCSV.foreach(filename, options) do |row|
    locations << Location.new(row.to_hash)
    #Breaking everything up helps with memory and improves speed
    if locations.size % chunk_size == 0
      #ar-extensions can create multiple database objects with a single statement through .import()
      Location.import locations, :validate => validate 
      locations = []
    end
  end
  Location.import locations, :validate => validate if locations.size > 0
end
 
load(ARGV[0], 1_000, false)
puts Location.count