ralphrodkey-kafka-rb 0.0.15

data/lib/kafka/message.rb

@@ -0,0 +1,209 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+
+  # A message. The format of a message is as follows:
+  #
+  # 4 byte big-endian int: length of message in bytes (including the rest of
+  #                        the header, but excluding the length field itself)
+  # 1 byte: "magic" identifier (format version number)
+  #
+  # If the magic byte == 0, there is one more header field:
+  #
+  # 4 byte big-endian int: CRC32 checksum of the payload
+  #
+  # If the magic byte == 1, there are two more header fields:
+  #
+  # 1 byte: "attributes" (flags for compression, codec etc)
+  # 4 byte big-endian int: CRC32 checksum of the payload
+  #
+  # All following bytes are the payload.
+  class Message
+
+    MAGIC_IDENTIFIER_DEFAULT = 0
+    MAGIC_IDENTIFIER_COMPRESSION = 1
+    NO_COMPRESSION = 0
+    GZIP_COMPRESSION = 1
+    SNAPPY_COMPRESSION = 2
+    BASIC_MESSAGE_HEADER = 'NC'.freeze
+    VERSION_0_HEADER = 'N'.freeze
+    VERSION_1_HEADER = 'CN'.freeze
+    COMPRESSION_CODEC_MASK = 0x03
+
+    attr_accessor :magic, :checksum, :payload
+
+    def initialize(payload = nil, magic = MAGIC_IDENTIFIER_DEFAULT, checksum = nil)
+      self.magic       = magic
+      self.payload     = payload || ""
+      self.checksum    = checksum || self.calculate_checksum
+      @compression = NO_COMPRESSION
+    end
+
+    def calculate_checksum
+      Zlib.crc32(self.payload)
+    end
+
+    def valid?
+      self.checksum == calculate_checksum
+    end
+
+    # Takes a byte string containing one or more messages; returns a MessageSet
+    # with the messages parsed from the string, and the number of bytes
+    # consumed from the string.
+    def self.parse_from(data)
+      messages = []
+      bytes_processed = 0
+
+      while bytes_processed <= data.length - 5 # 5 = size of BASIC_MESSAGE_HEADER
+        message_size, magic = data[bytes_processed, 5].unpack(BASIC_MESSAGE_HEADER)
+        break if bytes_processed + message_size + 4 > data.length # message is truncated
+
+        case magic
+        when MAGIC_IDENTIFIER_DEFAULT
+          # |  0  |  1  |  2  |  3  |  4  |  5  |  6  |  7  |  8  |  9      ...
+          # |                       |     |                       |
+          # |      message_size     |magic|        checksum       | payload ...
+          payload_size = message_size - 5 # 5 = sizeof(magic) + sizeof(checksum)
+          checksum = data[bytes_processed + 5, 4].unpack(VERSION_0_HEADER).shift
+          payload  = data[bytes_processed + 9, payload_size]
+          messages << Kafka::Message.new(payload, magic, checksum)
+
+        when MAGIC_IDENTIFIER_COMPRESSION
+          # |  0  |  1  |  2  |  3  |  4  |  5  |  6  |  7  |  8  |  9  | 10      ...
+          # |                       |     |     |                       |
+          # |         size          |magic|attrs|        checksum       | payload ...
+          payload_size = message_size - 6 # 6 = sizeof(magic) + sizeof(attrs) + sizeof(checksum)
+          attributes, checksum = data[bytes_processed + 5, 5].unpack(VERSION_1_HEADER)
+          payload = data[bytes_processed + 10, payload_size]
+
+          case attributes & COMPRESSION_CODEC_MASK
+          when NO_COMPRESSION # a single uncompressed message
+            messages << Kafka::Message.new(payload, magic, checksum)
+          when GZIP_COMPRESSION # a gzip-compressed message set -- parse recursively
+            uncompressed = Zlib::GzipReader.new(StringIO.new(payload)).read
+            message_set = parse_from(uncompressed)
+            raise 'malformed compressed message' if message_set.size != uncompressed.size
+            messages.concat(message_set.messages)
+          when SNAPPY_COMPRESSION # a snappy-compresses message set -- parse recursively
+            ensure_snappy! do
+              uncompressed = Snappy::Reader.new(StringIO.new(payload)).read
+              message_set = parse_from(uncompressed)
+              raise 'malformed compressed message' if message_set.size != uncompressed.size
+              messages.concat(message_set.messages)
+            end
+          else
+            # https://cwiki.apache.org/confluence/display/KAFKA/Compression
+            raise "Unsupported Kafka compression codec: #{attributes & COMPRESSION_CODEC_MASK}"
+          end
+
+        else
+          raise "Unsupported Kafka message version: magic number #{magic}"
+        end
+
+        bytes_processed += message_size + 4 # 4 = sizeof(message_size)
+      end
+
+      MessageSet.new(bytes_processed, messages)
+    end
+
+    def encode(compression = NO_COMPRESSION)
+      @compression = compression
+
+      self.payload = asciify_payload
+      self.payload = compress_payload if compression?
+
+      data = magic_and_compression + [calculate_checksum].pack("N") + payload
+      [data.length].pack("N") + data
+    end
+
+
+    # Encapsulates a list of Kafka messages (as Kafka::Message objects in the
+    # +messages+ attribute) and their total serialized size in bytes (the +size+
+    # attribute).
+    class MessageSet < Struct.new(:size, :messages); end
+
+    def self.ensure_snappy!
+      if Object.const_defined? "Snappy"
+        yield
+      else
+        fail "Snappy not available!"
+      end
+    end
+
+    def ensure_snappy! &block
+      self.class.ensure_snappy! &block
+    end
+
+    private
+
+    attr_reader :compression
+
+    def compression?
+      compression != NO_COMPRESSION
+    end
+
+    def magic_and_compression
+      if compression?
+        [MAGIC_IDENTIFIER_COMPRESSION, compression].pack("CC")
+      else
+        [MAGIC_IDENTIFIER_DEFAULT].pack("C")
+      end
+    end
+
+    def asciify_payload
+      if RUBY_VERSION[0, 3] == "1.8"
+        payload
+      else
+        payload.to_s.force_encoding(Encoding::ASCII_8BIT)
+      end
+    end
+
+    def compress_payload
+      case compression
+        when GZIP_COMPRESSION
+          gzip
+        when SNAPPY_COMPRESSION
+          snappy
+      end
+    end
+
+    def gzip
+      with_buffer do |buffer|
+        gz = Zlib::GzipWriter.new buffer, nil, nil
+        gz.write payload
+        gz.close
+      end
+    end
+
+    def snappy
+      ensure_snappy! do
+        with_buffer do |buffer|
+          Snappy::Writer.new buffer do |w|
+            w << payload
+          end
+        end
+      end
+    end
+
+    def with_buffer
+      buffer = StringIO.new
+      buffer.set_encoding Encoding::ASCII_8BIT unless RUBY_VERSION =~ /^1\.8/
+      yield buffer if block_given?
+      buffer.rewind
+      buffer.string
+    end
+  end
+end
+

data/lib/kafka/multi_producer.rb

@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+  class MultiProducer
+    include Kafka::IO
+
+    def initialize(options={})
+      self.host = options[:host] || HOST
+      self.port = options[:port] || PORT
+      self.compression = options[:compression] || Message::NO_COMPRESSION
+      self.connect(self.host, self.port)
+    end
+
+    def push(topic, messages, options={})
+      partition = options[:partition] || 0
+      self.write(Encoder.produce(topic, partition, messages, compression))
+    end
+
+    def multi_push(producer_requests)
+      self.write(Encoder.multiproduce(producer_requests, compression))
+    end
+  end
+end

data/lib/kafka/producer.rb

@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+  class Producer
+
+    include Kafka::IO
+
+    attr_accessor :topic, :partition
+
+    def initialize(options = {})
+      self.topic       = options[:topic]      || "test"
+      self.partition   = options[:partition]  || 0
+      self.host        = options[:host]       || HOST
+      self.port        = options[:port]       || PORT
+      self.compression = options[:compression] || Message::NO_COMPRESSION
+      self.connect(self.host, self.port)
+    end
+
+    def push(messages)
+      self.write(Encoder.produce(self.topic, self.partition, messages, compression))
+    end
+
+    def batch(&block)
+      batch = Kafka::Batch.new
+      block.call( batch )
+      push(batch.messages)
+      batch.messages.clear
+    end
+  end
+end

data/lib/kafka/producer_request.rb

@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Kafka
+  class ProducerRequest
+    attr_accessor :topic, :messages, :partition
+
+    def initialize(topic, messages, options={})
+      self.topic = topic
+      self.partition = options[:partition] || 0
+      self.messages = Array(messages)
+    end
+  end
+end

data/lib/kafka/request_type.rb

@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+  module RequestType
+    PRODUCE      = 0
+    FETCH        = 1
+    MULTIFETCH   = 2
+    MULTIPRODUCE = 3
+    OFFSETS      = 4
+  end
+end

data/lib/leffen-kafka.rb

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require File.join(File.dirname(__FILE__), "kafka")

data/lib/ralphrodkey-kafka.rb

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require File.join(File.dirname(__FILE__), "kafka")

data/spec/batch_spec.rb

@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+require File.dirname(__FILE__) + '/spec_helper'
+
+describe Batch do
+
+  before(:each) do
+    @batch = Batch.new
+  end
+
+  describe "batch messages" do
+    it "holds all messages to be sent" do
+      @batch.should respond_to(:messages)
+      @batch.messages.class.should eql(Array)
+    end
+
+    it "supports queueing/adding messages to be send" do
+      @batch.messages << mock(Kafka::Message.new("one"))
+      @batch.messages << mock(Kafka::Message.new("two"))
+      @batch.messages.length.should eql(2)
+    end
+  end
+end

data/spec/cli_spec.rb

@@ -0,0 +1,133 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+require File.dirname(__FILE__) + '/spec_helper'
+require 'kafka/cli'
+
+describe CLI do
+
+  before(:each) do
+    CLI.instance_variable_set("@config", {})
+    CLI.stub(:puts)
+  end
+
+  describe "should read from env" do
+    describe "kafka host" do
+      it "should read KAFKA_HOST from env" do
+        CLI.read_env("KAFKA_HOST" => "google.com")
+        CLI.config[:host].should == "google.com"
+      end
+
+      it "kafka port" do
+        CLI.read_env("KAFKA_PORT" => "1234")
+        CLI.config[:port].should == 1234
+      end
+
+      it "kafka topic" do
+        CLI.read_env("KAFKA_TOPIC" => "news")
+        CLI.config[:topic].should == "news"
+      end
+
+      it "kafka compression" do
+        CLI.read_env("KAFKA_COMPRESSION" => "no")
+        CLI.config[:compression].should == Message::NO_COMPRESSION
+
+        CLI.read_env("KAFKA_COMPRESSION" => "gzip")
+        CLI.config[:compression].should == Message::GZIP_COMPRESSION
+
+        CLI.read_env("KAFKA_COMPRESSION" => "snappy")
+        CLI.config[:compression].should == Message::SNAPPY_COMPRESSION
+      end
+    end
+  end
+
+  describe "should read from command line" do
+    it "kafka host" do
+      CLI.parse_args(%w(--host google.com))
+      CLI.config[:host].should == "google.com"
+
+      CLI.parse_args(%w(-h google.com))
+      CLI.config[:host].should == "google.com"
+    end
+
+    it "kafka port" do
+      CLI.parse_args(%w(--port 1234))
+      CLI.config[:port].should == 1234
+
+      CLI.parse_args(%w(-p 1234))
+      CLI.config[:port].should == 1234
+    end
+
+    it "kafka topic" do
+      CLI.parse_args(%w(--topic news))
+      CLI.config[:topic].should == "news"
+
+      CLI.parse_args(%w(-t news))
+      CLI.config[:topic].should == "news"
+    end
+
+    it "kafka compression" do
+      CLI.stub(:publish? => true)
+
+      CLI.parse_args(%w(--compression no))
+      CLI.config[:compression].should == Message::NO_COMPRESSION
+      CLI.parse_args(%w(-c no))
+      CLI.config[:compression].should == Message::NO_COMPRESSION
+
+      CLI.parse_args(%w(--compression gzip))
+      CLI.config[:compression].should == Message::GZIP_COMPRESSION
+      CLI.parse_args(%w(-c gzip))
+      CLI.config[:compression].should == Message::GZIP_COMPRESSION
+
+      CLI.parse_args(%w(--compression snappy))
+      CLI.config[:compression].should == Message::SNAPPY_COMPRESSION
+      CLI.parse_args(%w(-c snappy))
+      CLI.config[:compression].should == Message::SNAPPY_COMPRESSION
+    end
+
+    it "message" do
+      CLI.stub(:publish? => true)
+      CLI.parse_args(%w(--message YEAH))
+      CLI.config[:message].should == "YEAH"
+
+      CLI.parse_args(%w(-m YEAH))
+      CLI.config[:message].should == "YEAH"
+    end
+
+  end
+
+  describe "config validation" do
+    it "should assign a default port" do
+      CLI.stub(:exit)
+      CLI.stub(:puts)
+      CLI.validate_config
+      CLI.config[:port].should == Kafka::IO::PORT
+    end
+  end
+
+  it "should assign a default host" do
+    CLI.stub(:exit)
+    CLI.validate_config
+    CLI.config[:host].should == Kafka::IO::HOST
+  end
+
+
+  it "read compression method" do
+    CLI.string_to_compression("no").should == Message::NO_COMPRESSION
+    CLI.string_to_compression("gzip").should == Message::GZIP_COMPRESSION
+    CLI.string_to_compression("snappy").should == Message::SNAPPY_COMPRESSION
+    lambda { CLI.push(:string_to_compression,nil) }.should raise_error
+  end
+
+end