wukong-load 0.0.2 → 0.1.0

data/lib/wukong-load/loaders/kafka.rb

@@ -0,0 +1,98 @@
+require_relative('../loader')
+
+module Wukong
+  module Load
+
+    # Loads data into Kafka.
+    #
+    # Uses the `kafka-rb` gem to create a Kafka::Producer to write to
+    # Kafka.
+    #
+    # Allows loading records into a given topic on a given partition.
+    # Records can have fields `_topic` and `_partition` which override
+    # the given topic and partition on a per-record basis.
+    #
+    # The names of these fields within each record (`_topic` and
+    # `_partition`) can be customized.
+    class KafkaLoader < Loader
+
+      field :host,            String,  :default => 'localhost',  :doc => "Kafka broker host"
+      field :port,            Integer, :default => 9092,         :doc => "Kafka broker port"
+      field :topic,           String,  :default => 'test',       :doc => "Kafka topic"
+      field :topic_field,     String,  :default => '_topic',     :doc => "Field within records which names the Kafka topic"
+      field :partition,       Integer, :default => 0,            :doc => "Kafka partition"
+      field :partition_field, String,  :default => '_partition', :doc => "Field within records which names the Kafka partition"
+
+      description <<-EOF.gsub(/^ {8}/,'')
+        Loads newline-separated, JSON-formatted records over STDIN
+        into a Kafka queue.
+
+          $ cat data.json | wu-load kafka
+
+        By default, wu-load attempts to write each input record to a
+        local Kafka broker.
+
+        Input records will be written to a default Kafka topic on a
+        default partition.  Each record can have _topic and _partition
+        fields to override this on a per-record basis.
+
+        The fields used (_topic and _partition) can be changed:
+
+          $ cat data.json | wu-load kafka --host=10.123.123.123 --topic=hits --partition_field=segment_id
+      EOF
+
+      # The Kafka producer used to send messages to Kafka.
+      attr_accessor :producer
+
+      # Creates the producer.
+      def setup
+        begin
+          require 'kafka'
+        rescue => e
+          raise Error.new("Please ensure that the 'kafka-rb' gem is installed and available (in your Gemfile)")
+        end
+        log.debug("Connecting to Kafka broker at #{host}:#{port}...")
+        begin
+          self.producer = Kafka::MultiProducer.new(:host => host, :port => port)
+        rescue => e
+          raise Error.new(e.message)
+        end
+      end
+
+      # Load a single record into Kafka.
+      #
+      # @param [Hash] record
+      def load record
+        begin
+          topic     = topic_for(record)
+          partition = partition_for(record)
+          bytes     = producer.send(topic, messages_for(record), :partition => partition)
+          log.info("Wrote #{bytes} bytes to #{topic}/#{partition}")
+        rescue => e
+          handle_error(record, e)
+        end
+      end
+
+      # :nodoc:
+      def topic_for record
+        record[topic_field] || self.topic
+      end
+
+      # :nodoc:
+      def messages_for record
+        [Kafka::Message.new(MultiJson.dump(record))]
+      end
+
+      # :nodoc:
+      def partition_for record
+        record[partition_field] ? record[partition_field].to_i : partition
+      end
+      
+      register :kafka_loader
+      
+    end
+  end
+end
+
+    
+    

data/lib/wukong-load/loaders/mongodb.rb

@@ -0,0 +1,123 @@
+require_relative('../loader')
+
+module Wukong
+  module Load
+
+    # Loads data into MongoDB.
+    #
+    # Uses the 'mongo' gem to connect and write data.
+    #
+    # Allows loading records into a given database and collection.
+    # Records can have fields `_database` and `_collection` which
+    # override the given database and collection on a per-record
+    # basis.
+    #
+    # Records can have an `_id` field which indicates an update, not
+    # an insert.
+    #
+    # The names of these fields within each record (`_database`,
+    # `_collection`, and `_id`) can be customized.
+    class MongoDBLoader < Loader
+
+      field :host,             String, :default => 'localhost', :doc => "MongoDB host"
+      field :port,             Integer,:default => 27017, :doc => "Port on MongoDB host"
+      field :database,         String, :default => 'wukong', :doc => "Default MongoDB database"
+      field :collection,       String, :default => 'streaming_record', :doc => "Default MongoDB collection"
+      field :database_field,   String, :default => '_database', :doc => "Name of field in each record overriding default MongoDB database"
+      field :collection_field, String, :default => '_collection', :doc => "Name of field in each record overriding default MongoDB collection"
+      field :id_field,         String, :default => '_id', :doc => "Name of field in each record providing ID of existing MongoDB record to update"
+
+      description <<-EOF.gsub(/^ {8}/,'')
+        Loads newline-separated, JSON-formatted records over STDIN
+        into MongoDB.
+
+          $ cat data.json | wu-load mongodb
+
+        By default, wu-load attempts to write each input record to a
+        local MongoDB server.
+
+        Input records will be written to a default database and
+        collection.  Each record can have _database and _collection
+        fields to override this on a per-record basis.
+
+        Records with an _id field will be trigger updates, the rest
+        inserts.
+
+        All other fields within a record are assumed to be the names
+        of actual columns in the table.
+
+        The fields used (_index, _collection, and _id) can be changed:
+
+          $ cat data.json | wu-load mongodb --host=10.123.123.123 --database=web_events --collection=impressions --id_field=impression_id
+      EOF
+      
+      # The Mongo::MongoClient we'll use for talking to MongoDB.
+      attr_accessor :client
+
+      # Creates the client connection.
+      def setup
+        begin
+          require 'mongo'
+        rescue => e
+          raise Error.new("Please ensure that the 'mongo' gem is installed and available (in your Gemfile)")
+        end
+        h = host.gsub(%r{^http://},'')
+        log.debug("Connecting to MongoDB server at #{h}:#{port}...")
+        begin
+          self.client = Mongo::MongoClient.new(h, port)
+        rescue => e
+          raise Error.new(e.message)
+        end
+      end
+
+      # Load a single record into MongoDB.
+      #
+      # If the record has an ID, we'll issue an update, otherwise an
+      # insert.
+      #
+      # @param [record] Hash
+      def load record
+        id = id_for(record)
+        if id
+          res = collection_for(record).update({:_id => id}, record, :upsert => true)
+          if res['updatedExisting']
+            log.info("Updated #{id}")
+          else
+            log.info("Inserted #{id}")
+          end
+        else
+          res = collection_for(record).insert(record)
+          log.info("Inserted #{res}")
+        end
+      end
+
+      # :nodoc:
+      def database_for record
+        client[database_name_for(record)]
+      end
+
+      # :nodoc:
+      def collection_for record
+        database_for(record)[collection_name_for(record)]
+      end
+
+      # :nodoc:
+      def database_name_for record
+        record[database_field] || self.database
+      end
+
+      # :nodoc:
+      def collection_name_for record
+        record[collection_field] || self.collection
+      end
+
+      # :nodoc:
+      def id_for record
+        record[id_field]
+      end
+      
+      register :mongodb_loader
+      
+    end
+  end
+end

data/lib/wukong-load/loaders/sql.rb

@@ -0,0 +1,169 @@
+require_relative('../loader')
+
+module Wukong
+  module Load
+
+    # Loads data into SQL databases.
+    #
+    # Uses the 'mysql' gem to connect and write data.  Yes, MySQL !=
+    # SQL but we'll get there, I promise...
+    #
+    # Allows loading records into a given database and table.  Records
+    # can have fields `_database` and `_table` which override the
+    # given database and table on a per-record basis.
+    #
+    # Records can have an `_id` field which indicates an update, not
+    # an insert.
+    #
+    # The names of these fields within each record (`_database`,
+    # `_table`, and `_id`) can be customized.
+    class SQLLoader < Loader
+
+      field :host,             String, :default => 'localhost', :doc => "SQL host"
+      field :port,             Integer,:default => 3306, :doc => "Port on SQL host"
+      field :username,         String, :default => (ENV['USER'] || 'wukong'), :doc => "User to connect as"
+      field :password,         String, :doc => "Password for user"
+      field :database,         String, :default => 'wukong', :doc => "Default database"
+      field :table,            String, :default => 'streaming_record', :doc => "Default table"
+      field :database_field,   String, :default => '_database', :doc => "Name of field in each record overriding default database"
+      field :table_field,      String, :default => '_table', :doc => "Name of field in each record overriding default table"
+      field :id_field,         String, :default => '_id', :doc => "Name of field in each record providing ID of existing row to update"
+
+      description <<-EOF.gsub(/^ {8}/,'')
+        Loads newline-separated, JSON-formatted records over STDIN
+        into MySQL using its HTTP API.
+
+          $ cat data.json | wu-load sql
+
+        By default, wu-load attempts to write each input record to a
+        local SQL server.
+
+        Input records will be written to a default database and table.
+        Each record can have _database and _table fields to override
+        this on a per-record basis.
+
+        Records with an _id field will be trigger updates, the rest
+        inserts.
+
+        All other fields within a record are assumed to be the names
+        of actual columns in the table.
+
+        The fields used (_index, _table, and _id) can be changed:
+
+          $ cat data.json | wu-load sql --host=10.123.123.123 --database=web_events --table=impressions --id_field=impression_id
+      EOF
+      
+      # The Mongo::MongoClient we'll use for talking to MongoDB.
+      attr_accessor :client
+
+      # Creates the client connection.
+      def setup
+        begin
+          require 'mysql2'
+        rescue => e
+          raise Error.new("Please ensure that the 'mysql2' gem is installed and available (in your Gemfile)")
+        end
+        log.debug("Connecting to SQL server at <#{host}:#{port}> as <#{username}>#{' using password' if password}...")
+        begin
+          self.client = Mysql2::Client.new(sql_params)
+        rescue => e
+          raise Error.new(e)
+        end
+      end
+
+      # :nodoc:
+      def sql_params
+        {:host => host, :port => port}.tap do |params|
+          params[:username] if username
+          params[:password] if password
+        end
+      end
+
+      # Load a single record into the database.
+      #
+      # If the record has an ID, we'll issue an update, otherwise an
+      # insert.
+      #
+      # @param [record] Hash
+      def load record
+        id = id_for(record)
+        if id
+          perform_query(update_query(record))
+          log.info("Updated #{id}")
+        else
+          perform_query(insert_query(record))
+          log.info("Inserted")
+        end
+      end
+      
+      # :nodoc:
+      def insert_query record
+        "INSERT INTO #{database_name_for(record)}.#{table_name_for(record)} (#{fields_of(record)}) VALUES (#{values_of(record)}) ON DUPLICATE KEY UPDATE #{fields_and_values_of(record)}"
+      end
+
+      # :nodoc:
+      def update_query record
+        "UPDATE #{database_name_for(record)}.#{table_name_for(record)} SET #{fields_and_values_of(record)} WHERE `id`=#{id_for(record)}"
+      end
+
+      # :nodoc:
+      def field_names_of record
+        record.keys.reject { |key| [database_field, table_field, id_field].include?(key) }.sort
+      end
+
+      # :nodoc:
+      def fields_of record
+        field_names_of(record).map { |name| identifier_for(name) }.join(', ')
+      end
+
+      # :nodoc:
+      def values_of record
+        field_names_of(record).map { |name| value_for(record[name]) }.join(', ')
+      end
+
+      # :nodoc:
+      def fields_and_values_of record
+        field_names_of(record).map { |name| [identifier_for(name), value_for(record[name])].join('=') }.join(', ')
+      end
+
+      # :nodoc:
+      def database_name_for record
+        identifier_for(record[database_field] || self.database)
+      end
+
+      # :nodoc:
+      def table_name_for record
+        identifier_for(record[table_field] || self.table)
+      end
+
+      # :nodoc:
+      def identifier_for thing
+        '`' + client.escape(thing.to_s) + '`'
+      end
+
+      # :nodoc:
+      def value_for thing
+        case thing
+        when Fixnum then thing
+        when nil    then 'NULL'
+        else
+          '"' + client.escape(thing.to_s) + '"'
+        end
+      end
+      
+
+      # :nodoc:
+      def id_for record
+        value_for(record[id_field]) if record[id_field]
+      end
+
+      # :nodoc:
+      def perform_query query
+        client.query query
+      end
+      
+      register :sql_loader
+      
+    end
+  end
+end

data/lib/wukong-load/models/http_request.rb

@@ -0,0 +1,60 @@
+module Wukong
+
+  # Represents a generic HTTP request.
+  class HttpRequest
+    
+    include Gorillib::Model
+
+    field :timestamp,  Integer, :doc => "Timestamp at which the HTTP request was received"
+    field :verb,       String,  :doc => "HTTP verb of the request"
+    field :path,       String,  :doc => "Absolute path to the resource requested"
+    field :params,     Hash,    :doc => "Query parameters contained in the request", :default => {}
+    field :headers,    Hash,    :doc => "HTTP headers of the request", :default => {}
+    field :ip_address, String,  :doc => "IP address of the client"
+    field :body,       String,  :doc => "Body of the request"
+
+    # Return the URL of this request.
+    #
+    # @return [String]
+    def url
+      File.join(headers['Host'] || '', (path || ''))
+    end
+
+    # Return the HTTP Referer of this request.
+    #
+    # @return [String]
+    def referer
+      headers['Referer']
+    end
+    alias_method :referrer, :referer
+
+    # Return the HTTP User-Agent of this request.
+    #
+    # @return [String]
+    def user_agent
+      headers['User-Agent']
+    end
+
+    # Return the HTTP Cookie of this request.
+    #
+    # @return [String]
+    def cookie
+      headers['Cookie']
+    end
+
+    # Return the "best" IP address from this request.
+    #
+    # Will return the first IP address in the HTTP X-Forwarded-For chain
+    # if present, otherwise will return the IP address of the request
+    # itself.
+    #
+    # @return [String]
+    def best_ip_address
+      ip_string = headers['X-Forwarded-For']
+      return ip_address if ip_string.blank?
+      ips = ip_string.split(/\s*,\s*/)
+      ips.empty? ? ip_address : ips.first # client comes first, then proxies in order
+    end
+    
+  end
+end