Fluentdのout_geoipをちょっと改造したやつ

out_geoip_custom.rb

require 'fluent/mixin/rewrite_tag_name'

class Fluent::GeoipCustomOutput < Fluent::BufferedOutput
  Fluent::Plugin.register_output('geoip_custom', self)

  config_param :geoip_database, :string, :default => File.dirname(__FILE__) + '/../../../data/GeoLiteCity.dat'
  config_param :geoip_lookup_key, :string, :default => 'host'
  config_param :tag, :string, :default => nil
  config_param :skip_adding_null_record, :bool, :default => false
  config_param :result_key_suffix, :string, :default => '_geoip'

  include Fluent::HandleTagNameMixin
  include Fluent::SetTagKeyMixin
  config_set_default :include_tag_key, false

  include Fluent::Mixin::RewriteTagName
  config_param :hostname_command, :string, :default => 'hostname'

  config_param :flush_interval, :time, :default => 0
  config_param :log_level, :string, :default => 'warn'

  # Define `log` method for v0.10.42 or earlier
  unless method_defined?(:log)
    define_method("log") { $log }
  end

  def initialize
    require 'geoip'
    require 'yajl'

    super
  end

  def configure(conf)
    super

    @geoip_lookup_key = @geoip_lookup_key.split(/\s*,\s*/)

    if ( !@tag && !@remove_tag_prefix && !@remove_tag_suffix && !@add_tag_prefix && !@add_tag_suffix )
      raise Fluent::ConfigError, "geoip: required at least one option of 'tag', 'remove_tag_prefix', 'remove_tag_suffix', 'add_tag_prefix', 'add_tag_suffix'."
    end

    @geoip = GeoIP::City.new(@geoip_database, :memory, false)
  end

  def start
    super
  end

  def format(tag, time, record)
    [tag, time, record].to_msgpack
  end

  def shutdown
    super
  end

  def write(chunk)
    chunk.msgpack_each do |tag, time, record|
      Fluent::Engine.emit(tag, time, add_geoip_field(record))
    end
  end

  private

  def add_geoip_field(record)
    geodata = get_geodata(record)
    geodata.each do |key, value|
      record.store(key + @result_key_suffix, value)
    end

    return record
  end

  def get_geodata(record)
    geodata = {}

    @geoip_lookup_key.each do |field|
      if record[field].nil? then
        next
      end
      geos = []
      ips = record[field].split(/\s*,\s*/)
      ips.each do |ip|
        geos << @geoip.look_up(ip)
      end
      geodata.store(field, geos)
    end

    return geodata
  end
end