box-o-sand/redis/with-ruby/earthquakes.rb

140 lines
3.6 KiB
Ruby

require 'csv'
require 'redis/connection/hiredis'
require 'redis'
require 'logger'
require 'open-uri'
class Earthquakes
attr_accessor :id_set, :region_set, :magnitude_set
attr_accessor :magnitude_scores_set, :data_url
def initialize
@redis = Redis.new
@id_set = 'earthquake-ids'
@region_set = 'earthquake-regions'
@magnitude_set = 'earthquake-magnitudes'
@magnitude_scores_set = 'earthquakes-magnitudes-scores'
@data_url =
'http://earthquake.usgs.gov/earthquakes/catalogs/eqs7day-M1.txt'
@logger = Logger.new(STDOUT)
@logger.formatter = proc do |severity, datetime, progname, msg|
"#{datetime}: #{msg}\n"
end
end
def get_info(earthquake_id)
@redis.hgetall(with_prefix(earthquake_id, 'eqid'))
end
def in_region(region)
@redis.smembers(with_prefix(region, 'eqregion'))
end
def with_magnitude_in_range(min_magnitude, max_magnitude)
min_mag, max_mag =
magnitude_as_int(min_magnitude), magnitude_as_int(max_magnitude)
@logger.debug(
"Fetching magnitudes with scores in range #{min_mag}-#{max_mag}")
@redis.zrange(@magnitude_scores_set, min_mag, max_mag)
end
def all_ids
@redis.smembers(@id_set).map { |i| sans_prefix(i, 'eqid') }
end
def all_regions
@redis.smembers(@region_set).map { |r| sans_prefix(r, 'eqregion') }
end
def all_magnitudes
@redis.smembers(@magnitude_set).map { |m| sans_prefix(m, 'eqmagnitude') }
end
def drop_all
@logger.info('Removing data for each earthquake')
all_ids.each do |i|
@logger.info("Removing data for earthquake #{sans_prefix(i, 'eqid')}")
@redis.del(i)
end
@logger.info('Removing earthquake id set')
@redis.del(@id_set)
@logger.info('Removing earthquake region set')
@redis.del(@region_set)
@logger.info('Removing earthquake magnitude set')
@redis.del(@magnitude_set)
@logger.info('Removing earthquake magnitude scores set')
@redis.del(@magnitude_scores_set)
end
def load_data
data = CSV.parse(
open(@data_url).read, :headers => true, :header_converters => [:symbol]
)
data.each do |row|
@logger.info("Adding earthquake #{row[:eqid]}")
load_one_row(row)
end
end
private
def load_one_row(row)
@redis.multi do |redis_multi|
id_key = with_prefix(row[:eqid].downcase, 'eqid')
load_id_and_attrs_for_row(row, id_key, redis_multi)
load_region_for_row(row, id_key, redis_multi)
load_magnitude_for_row(row, id_key, redis_multi)
end
end
def load_id_and_attrs_for_row(row, id_key, redis_multi)
redis_multi.sadd(@id_set, id_key)
row.headers.each do |header|
redis_multi.hset(id_key, header.downcase, row[header].downcase)
end
end
def load_region_for_row(row, id_key, redis_multi)
region_key = with_prefix(row[:region].downcase, 'eqregion')
redis_multi.sadd(region_key, id_key)
redis_multi.sadd(@region_set, region_key)
end
def load_magnitude_for_row(row, id_key, redis_multi)
magnitude_key = with_prefix(row[:magnitude], 'eqmagnitude')
redis_multi.sadd(magnitude_key, id_key)
redis_multi.sadd(@magnitude_set, magnitude_key)
score = magnitude_as_int(row[:magnitude])
@logger.info("Setting score of #{score} for row #{row.inspect}")
redis_multi.zadd(@magnitude_scores_set, score, id_key)
end
def with_prefix(key, prefix)
if not /^#{prefix}:/.match(key)
return "#{prefix}:#{key}"
end
key
end
def sans_prefix(key, prefix)
key.gsub(/^#{prefix}:/, '')
end
def magnitude_as_int(magnitude)
magnitude.to_s.gsub(/\./, '').to_i
end
end
if $0 == __FILE__
Earthquakes.new.send(ARGV.first)
end