wukong-load 0.0.2 → 0.1.0

data/lib/wukong-load/source_driver.rb

@@ -0,0 +1,46 @@
+module Wukong
+  module Load
+    class SourceDriver < Wukong::Local::StdioDriver
+      include Logging
+
+      attr_accessor :index, :batch_size
+
+      def post_init
+        super()
+        self.index = 1
+        self.batch_size = settings[:batch_size].to_i if settings[:batch_size] && settings[:batch_size].to_i > 0
+      end
+
+      def self.start(label, settings={})
+        driver = new(:foobar, label, settings)
+        driver.post_init
+
+        period = case
+        when settings[:period]  then settings[:period]
+        when settings[:per_sec] then (1.0 / settings[:per_sec]) rescue 1.0
+        else 1.0
+        end
+        driver.create_event
+        EventMachine::PeriodicTimer.new(period) { driver.create_event }
+      end
+
+      def create_event
+        receive_line(index.to_s)
+        self.index += 1
+        finalize_dataflow if self.batch_size && (self.index % self.batch_size) == 0
+      end
+
+      # :nodoc:
+      #
+      # Not sure why I have to add the call to $stdout.flush at the
+      # end of this method.  Supposedly $stdout.sync is called during
+      # the #setup method in StdoutProcessor in
+      # wukong/widget/processors.  Doesn't that do this?
+      def process record
+        $stdout.puts record
+        $stdout.flush
+      end
+      
+    end
+  end
+end

data/lib/wukong-load/source_runner.rb

@@ -0,0 +1,36 @@
+require_relative('source_driver')
+module Wukong
+  module Load
+
+    # Runs the wu-source command.
+    class SourceRunner < Wukong::Local::LocalRunner
+
+      usage "PROCESSOR|DATAFLOW"
+
+      description <<-EOF.gsub(/^ {8}/,'')
+
+        wu-source is a tool for using Wukong processors as sources of
+        data in streams.
+
+        Run any Wukong processor as a source for data:
+
+          $ wu-source fake_log_data
+          205.4.75.208 - 3918471017 [27/Nov/2012:05:06:57 -0600] "GET /products/eget HTTP/1.0" 200 25600
+          63.181.105.15 - 3650805763 [27/Nov/2012:05:06:57 -0600] "GET /products/lacinia-nulla-vitae HTTP/1.0" 200 3790
+          227.190.78.101 - 39543891 [27/Nov/2012:05:06:58 -0600] "GET /products/odio-nulla-nulla-ipsum HTTP/1.0" 200 31718
+          ...
+
+        The fake_log_data processor will receive an event once every
+        second.  Each event will consist of a single string giving a
+        consecutive integer starting with '1' as the first event.
+      EOF
+      
+      include Logging
+
+      def driver
+        SourceDriver
+      end
+
+    end
+  end
+end

data/lib/wukong-load/version.rb

@@ -1,6 +1,6 @@
 module Wukong
   module Load
     # The current version of Wukong-Load
-    VERSION = '0.0.2'
+    VERSION = '0.1.0'
   end
 end

data/spec/spec_helper.rb

@@ -3,5 +3,18 @@ require 'wukong/spec_helpers'
 
 RSpec.configure do |config|
   config.mock_with :rspec
+  
   include Wukong::SpecHelpers
+
+  config.before(:each) do
+    Wukong::Log.level = Log4r::OFF
+  end
+  
+  def root
+    @root ||= Pathname.new(File.expand_path('../..', __FILE__))
+  end
+
+  def load_runner *args, &block
+    runner(Wukong::Load::LoadRunner, 'wu-load', *args)
+  end
 end

data/spec/wukong-load/loaders/elasticsearch_spec.rb

@@ -0,0 +1,142 @@
+require 'spec_helper'
+
+describe Wukong::Load::ElasticsearchLoader do
+
+  let(:loader)                   { Wukong::Load::ElasticsearchLoader.new                                }
+  let(:loader_with_custom_index) { Wukong::Load::ElasticsearchLoader.new(:index    => 'custom_index')   }
+  let(:loader_with_custom_type)  { Wukong::Load::ElasticsearchLoader.new(:es_type  => 'custom_es_type') }
+  let(:loader_with_custom_id)    { Wukong::Load::ElasticsearchLoader.new(:id_field => '_custom_id')     }
+
+  let(:record)                     { {'text' => 'hi'                                        } }
+  let(:record_with_index)          { {'text' => 'hi', '_index'          => 'custom_index'   } }
+  let(:record_with_custom_index)   { {'text' => 'hi', '_custom_index'   => 'custom_index'   } }
+  let(:record_with_es_type)        { {'text' => 'hi', '_es_type'        => 'custom_es_type' } }
+  let(:record_with_custom_es_type) { {'text' => 'hi', '_custom_es_type' => 'custom_es_type' } }
+  let(:record_with_id)             { {'text' => 'hi', '_id'             => 'the_id'         } }
+  let(:record_with_custom_id)      { {'text' => 'hi', '_custom_id'      => 'the_id'         } }
+
+  it_behaves_like 'a processor', :named => :elasticsearch_loader
+
+  context "without an Elasticsearch available" do
+    before do
+      Net::HTTP.should_receive(:new).and_raise(StandardError)
+    end
+
+    it "raises an error on setup" do
+      expect { processor(:elasticsearch_loader) }.to raise_error(Wukong::Error)
+    end
+  end
+
+  context "routes" do
+    context "all records" do
+      it "to a default index" do
+        loader.index_for(record).should == loader.index
+      end
+      it "to a given index" do
+        loader_with_custom_index.index_for(record).should == 'custom_index'
+      end
+      it "to a default type" do
+        loader.es_type_for(record).should == loader.es_type
+      end
+      it "to a given type" do
+        loader_with_custom_type.es_type_for(record).should == 'custom_es_type'
+      end
+    end
+    
+    context "records having a value for" do
+      it "default index field to the given index" do
+        loader.index_for(record_with_index).should == 'custom_index'
+      end
+      it "given index field to the given index" do
+        loader_with_custom_index.index_for(record_with_custom_index).should == 'custom_index'
+      end
+      it "default type field to the given type" do
+        loader.es_type_for(record_with_es_type).should == 'custom_es_type'
+      end
+      it "given type field to the given type" do
+        loader_with_custom_type.es_type_for(record_with_custom_es_type).should == 'custom_es_type'
+      end
+    end
+  end
+
+  context "detects IDs" do
+    it "based on the absence of a default ID field" do
+      loader.id_for(record).should be_nil
+    end
+    it "based on the value of a default ID field" do
+      loader.id_for(record_with_id).should == 'the_id'
+    end
+    it "based on the value of a custom ID field" do
+      loader_with_custom_id.id_for(record_with_custom_id).should == 'the_id'
+    end
+  end
+
+  context "sends" do
+    it "create requests on a record without an ID" do
+      loader.should_receive(:request).with(Net::HTTP::Post, '/foo/bar', kind_of(Hash))
+      loader.load({'_index' => 'foo', '_es_type' => 'bar'})
+    end
+    
+    it "update requests on a record with an ID" do
+      processor(:elasticsearch_loader) do |proc|
+        proc.should_receive(:request).with(Net::HTTP::Put, '/foo/bar/1', kind_of(Hash))
+        proc.load({'_index' => 'foo', '_es_type' => 'bar', '_id' => '1'})
+      end
+    end
+  end
+
+  context "receives" do
+    let(:connection) { double()                                              }
+    before           { Net::HTTP.should_receive(:new).and_return(connection) }
+    
+    let(:ok) do
+      mock("Net::HTTPOK").tap do |response|
+        response.stub!(:code).and_return('200')
+        response.stub!(:body).and_return('{"ok": true}')
+      end
+    end
+    let(:created) do
+      mock("Net::HTTPCreated").tap do |response|
+        response.stub!(:code).and_return('201')
+        response.stub!(:body).and_return('{"created": true}')
+      end
+    end
+    let(:not_found) do
+      mock("Net::HTTPNotFound").tap do |response|
+        response.stub!(:code).and_return('404')
+        response.stub!(:body).and_return('{"error": "Not found"}')
+      end
+    end
+    
+    context "201 Created" do
+      before { connection.should_receive(:request).with(kind_of(Net::HTTP::Post)).and_return(created) }
+      it "by logging an INFO message" do
+        processor(:elasticsearch_loader) do |proc|
+          proc.log.should_receive(:info)
+          proc.load(record)
+        end
+      end
+    end
+
+    context "200 OK" do
+      before { connection.should_receive(:request).with(kind_of(Net::HTTP::Put)).and_return(ok) }
+      it "by logging an INFO message" do
+        processor(:elasticsearch_loader) do |proc|
+          proc.log.should_receive(:info)
+          proc.load(record_with_id)
+        end
+      end
+    end
+
+    context "an error response from Elasticsearch" do
+      before { connection.should_receive(:request).with(kind_of(Net::HTTP::Post)).and_return(not_found) }
+      it "by logging an ERROR message" do
+        processor(:elasticsearch_loader) do |proc|
+          proc.log.should_receive(:error)
+          proc.load(record)
+        end
+      end
+    end
+    
+  end
+end

data/spec/wukong-load/loaders/kafka_spec.rb

@@ -0,0 +1,72 @@
+require 'spec_helper'
+require 'kafka'
+
+describe Wukong::Load::KafkaLoader do
+
+  let(:loader)                             { Wukong::Load::KafkaLoader.new                                           }
+  let(:loader_with_custom_topic)           { Wukong::Load::KafkaLoader.new(:topic           => 'custom'            ) }
+  let(:loader_with_custom_topic_field)     { Wukong::Load::KafkaLoader.new(:topic_field     => '_custom_topic'     ) }
+  let(:loader_with_custom_partition)       { Wukong::Load::KafkaLoader.new(:partition       => 1                   ) }
+  let(:loader_with_custom_partition_field) { Wukong::Load::KafkaLoader.new(:partition_field => '_custom_partition' ) }
+
+  let(:record)                             { {'text' => 'hi'                                  } }
+  let(:record_with_topic_field)            { {'text' => 'hi', '_topic'            => 'custom' } }
+  let(:record_with_custom_topic_field)     { {'text' => 'hi', '_custom_topic'     => 'custom' } }
+  let(:record_with_partition_field)        { {'text' => 'hi', '_partition'        => 1        } }
+  let(:record_with_custom_partition_field) { {'text' => 'hi', '_custom_partition' => 1        } }
+
+
+  it "raises an error on setup if it can't connect to Kafka" do
+    Kafka::MultiProducer.should_receive(:new).and_raise(StandardError)
+    expect { processor(:kafka_loader) }.to raise_error(Wukong::Error)
+  end
+  
+  context "with a Kafka available" do
+    before do
+      @producer = double()
+      Kafka::MultiProducer.stub!(:new).and_return(@producer)
+    end
+    it_behaves_like 'a processor', :named => :kafka_loader
+
+    it "produces an INFO log message on every write" do
+      @producer.should_receive(:send)
+      processor(:kafka_loader) do |proc|
+        proc.log.should_receive(:info)
+        proc.load(record)
+      end
+    end
+    
+  end
+  
+  context "routes" do
+    context "all records" do
+      it "to a default topic" do
+        loader.topic_for(record).should == loader.topic
+      end
+      it "to a given topic" do
+        loader_with_custom_topic.topic_for(record).should == 'custom'
+      end
+      it "to a default partition" do
+        loader.partition_for(record).should == loader.partition
+      end
+      it "to a given partition" do
+        loader_with_custom_partition.partition_for(record).should == 1
+      end
+    end
+    context "records having a value for" do
+      it "default topic field to the given topic" do
+        loader.topic_for(record_with_topic_field).should == 'custom'
+      end
+      it "given topic field to the given topic" do
+        loader_with_custom_topic.topic_for(record_with_topic_field).should == 'custom'
+      end
+      it "default partition field to the given partition" do
+        loader.partition_for(record_with_partition_field).should == 1
+      end
+      it "given partition field to the given partition" do
+        loader_with_custom_partition.partition_for(record_with_partition_field).should == 1
+      end
+    end
+  end
+
+end

data/spec/wukong-load/loaders/mongodb_spec.rb

@@ -0,0 +1,100 @@
+require 'spec_helper'
+require 'mongo'
+
+describe Wukong::Load::MongoDBLoader do
+
+  let(:loader)                   { Wukong::Load::MongoDBLoader.new                                }
+  let(:loader_with_custom_database) { Wukong::Load::MongoDBLoader.new(:database    => 'custom_database')   }
+  let(:loader_with_custom_collection)  { Wukong::Load::MongoDBLoader.new(:collection  => 'custom_collection') }
+  let(:loader_with_custom_id)    { Wukong::Load::MongoDBLoader.new(:id_field => '_custom_id')     }
+
+  let(:record)                     { {'text' => 'hi'                                        } }
+  let(:record_with_database)          { {'text' => 'hi', '_database'          => 'custom_database'   } }
+  let(:record_with_custom_database)   { {'text' => 'hi', '_custom_database'   => 'custom_database'   } }
+  let(:record_with_collection)        { {'text' => 'hi', '_collection'        => 'custom_collection' } }
+  let(:record_with_custom_collection) { {'text' => 'hi', '_custom_collection' => 'custom_collection' } }
+  let(:record_with_id)             { {'text' => 'hi', '_id'             => 'the_id'         } }
+  let(:record_with_custom_id)      { {'text' => 'hi', '_custom_id'      => 'the_id'         } }
+
+  
+
+  context "without an MongoDB available" do
+    before do
+      Mongo::MongoClient.should_receive(:new).and_raise(StandardError)
+    end
+
+    it "raises an error on setup" do
+      expect { processor(:mongodb_loader) }.to raise_error(Wukong::Error)
+    end
+  end
+
+  context "with a MongoDB available" do
+    before do
+      @client = double()
+      Mongo::MongoClient.stub!(:new).and_return(@client)
+    end
+    
+    it_behaves_like 'a processor', :named => :mongodb_loader
+    
+    context "routes" do
+      context "all records" do
+        it "to a default database" do
+          loader.database_name_for(record).should == loader.database
+        end
+        it "to a given database" do
+          loader_with_custom_database.database_name_for(record).should == 'custom_database'
+        end
+        it "to a default collection" do
+          loader.collection_name_for(record).should == loader.collection
+        end
+        it "to a given collection" do
+          loader_with_custom_collection.collection_name_for(record).should == 'custom_collection'
+        end
+      end
+      
+      context "records having a value for" do
+        it "default database field to the given database" do
+          loader.database_name_for(record_with_database).should == 'custom_database'
+        end
+        it "given database field to the given database" do
+          loader_with_custom_database.database_name_for(record_with_custom_database).should == 'custom_database'
+        end
+        it "default collection field to the given collection" do
+          loader.collection_name_for(record_with_collection).should == 'custom_collection'
+        end
+        it "given collection field to the given collection" do
+          loader_with_custom_collection.collection_name_for(record_with_custom_collection).should == 'custom_collection'
+        end
+      end
+    end
+
+    context "detects IDs" do
+      it "based on the absence of a default ID field" do
+        loader.id_for(record).should be_nil
+      end
+      it "based on the value of a default ID field" do
+        loader.id_for(record_with_id).should == 'the_id'
+      end
+      it "based on the value of a custom ID field" do
+        loader_with_custom_id.id_for(record_with_custom_id).should == 'the_id'
+      end
+    end
+
+    context "sends" do
+      before do
+        @collection = double()
+        loader.stub!(:collection_for).and_return(@collection)
+      end
+      it "insert requests on a record without an ID" do
+        @collection.should_receive(:insert).with(kind_of(Hash))
+        loader.load({'_database' => 'foo', '_collection' => 'bar'})
+      end
+      
+      it "update requests on a record with an ID" do
+        @collection.should_receive(:update).with({:_id => '1'}, kind_of(Hash), :upsert => true).and_return({})
+        loader.load({'_database' => 'foo', '_collection' => 'bar', '_id' => '1'})
+      end
+    end
+    
+  end
+end

data/spec/wukong-load/loaders/sql_spec.rb

@@ -0,0 +1,112 @@
+require 'spec_helper'
+require 'mongo'
+require 'mysql2'
+
+describe Wukong::Load::SQLLoader do
+
+  let(:loader)                   { Wukong::Load::SQLLoader.new                                }
+  let(:loader_with_custom_database) { Wukong::Load::SQLLoader.new(:database    => 'custom_database')   }
+  let(:loader_with_custom_table)  { Wukong::Load::SQLLoader.new(:table  => 'custom_table') }
+  let(:loader_with_custom_id)    { Wukong::Load::SQLLoader.new(:id_field => '_custom_id')     }
+
+  let(:record)                     { {'text' => 'hi'                                        } }
+  let(:record_with_database)          { {'text' => 'hi', '_database'          => 'custom_database'   } }
+  let(:record_with_custom_database)   { {'text' => 'hi', '_custom_database'   => 'custom_database'   } }
+  let(:record_with_table)        { {'text' => 'hi', '_table'        => 'custom_table' } }
+  let(:record_with_custom_table) { {'text' => 'hi', '_custom_table' => 'custom_table' } }
+  let(:record_with_id)             { {'text' => 'hi', '_id'             => 'the_id'         } }
+  let(:record_with_custom_id)      { {'text' => 'hi', '_custom_id'      => 'the_id'         } }
+
+  
+
+  context "without an SQL available" do
+    before do
+      Mysql2::Client.should_receive(:new).and_raise(StandardError)
+    end
+
+    it "raises an error on setup" do
+      expect { processor(:sql_loader) }.to raise_error(Wukong::Error)
+    end
+  end
+
+  context "with a SQL available" do
+    before do
+      @client = double()
+      def @client.escape record ; record.to_s ; end
+      Mysql2::Client.stub!(:new).and_return(@client)
+    end
+    
+    it_behaves_like 'a processor', :named => :sql_loader
+    
+    context "routes" do
+      context "all records" do
+        it "to a default database" do
+          loader.setup
+          loader.database_name_for(record).should == '`wukong`'
+        end
+        it "to a given database" do
+          loader_with_custom_database.setup
+          loader_with_custom_database.database_name_for(record).should == '`custom_database`'
+        end
+        it "to a default table" do
+          loader.setup
+          loader.table_name_for(record).should == '`streaming_record`'
+        end
+        it "to a given table" do
+          loader_with_custom_table.setup
+          loader_with_custom_table.table_name_for(record).should == '`custom_table`'
+        end
+      end
+      
+      context "records having a value for" do
+        it "default database field to the given database" do
+          loader.setup
+          loader.database_name_for(record_with_database).should == '`custom_database`'
+        end
+        it "given database field to the given database" do
+          loader_with_custom_database.setup
+          loader_with_custom_database.database_name_for(record_with_custom_database).should == '`custom_database`'
+        end
+        it "default table field to the given table" do
+          loader.setup
+          loader.table_name_for(record_with_table).should == '`custom_table`'
+        end
+        it "given table field to the given table" do
+          loader_with_custom_table.setup
+          loader_with_custom_table.table_name_for(record_with_custom_table).should == '`custom_table`'
+        end
+      end
+    end
+
+    context "detects IDs" do
+      it "based on the absence of a default ID field" do
+        loader.setup
+        loader.id_for(record).should be_nil
+      end
+      it "based on the value of a default ID field" do
+        loader.setup
+        loader.id_for(record_with_id).should == '"the_id"'
+      end
+      it "based on the value of a custom ID field" do
+        loader_with_custom_id.setup
+        loader_with_custom_id.id_for(record_with_custom_id).should == '"the_id"'
+      end
+    end
+
+    context "sends" do
+      before do
+        loader.setup
+      end
+      it "insert requests on a record without an ID" do
+        @client.should_receive(:query).with(%Q{INSERT INTO `foo`.`bar` (`age`, `email`, `name`) VALUES (58, "jerry@nbc.com", "Jerry Seinfeld") ON DUPLICATE KEY UPDATE `age`=58, `email`="jerry@nbc.com", `name`="Jerry Seinfeld"})
+        loader.load({'_database' => 'foo', '_table' => 'bar', 'name' => 'Jerry Seinfeld', 'email' => 'jerry@nbc.com', 'age' => 58})
+      end
+      
+      it "update requests on a record with an ID" do
+        @client.should_receive(:query).with(%Q{UPDATE `foo`.`bar` SET `age`=58, `email`="jerry@nbc.com", `name`="Jerry Seinfeld" WHERE `id`="1"})
+        loader.load({'_database' => 'foo', '_table' => 'bar', '_id' => '1', 'name' => 'Jerry Seinfeld', 'email' => 'jerry@nbc.com', 'age' => 58})
+      end
+    end
+    
+  end
+end