genki-kyototycoon 0.6.0.1

data/lib/kyototycoon/stream.rb

@@ -0,0 +1,74 @@
+# -- coding: utf-8
+
+=begin
+$ cat foo.rb
+require "rubygems"
+require "kyototycoon"
+
+KyotoTycoon::Stream.run($stdin) do |line|
+  ... do some stuff ..
+end
+
+$ ktremotemgr slave -uw | ruby foo.rb
+=end
+
+class KyotoTycoon
+  module Stream
+    def self.run(io=$stdin, &block)
+      io.each_line{|line|
+        line = Line.new(*line.strip.split("\t", 5))
+        block.call(line)
+      }
+    end
+
+    class Line < Struct.new(:ts, :sid, :db, :cmd, :raw_args)
+      def args
+        @args ||= begin
+          return [] if raw_args.nil?
+          k,v = *raw_args.split("\t").map{|v| v.unpack('m').first}
+          return [k] unless v
+          xt = 0
+          v.unpack('C5').each{|num|
+            xt = (xt << 8) +  num
+          }
+          v = v[5, v.length]
+          [k, v, xt]
+        end
+      end
+
+      def key
+        @key ||= begin
+          args.first || nil
+        end
+      end
+
+      def value
+        @value ||= begin
+          args[1] || nil
+        end
+      end
+
+      def xt
+        @xt ||= begin
+          args[2] || nil
+        end
+      end
+
+      def xt_time
+        @xt_time ||= begin
+          if args[2]
+            # if not set xt:
+            # Time.at(1099511627775) # => 36812-02-20 09:36:15 +0900
+            Time.at(args[2].to_i) 
+          else
+            Time.at(0)
+          end
+        end
+      end
+      
+      def time
+        @time ||= Time.at(*[ts[0,10], ts[10, ts.length]].map(&:to_i))
+      end
+    end
+  end
+end

data/lib/kyototycoon/tsvrpc/skinny.rb

@@ -0,0 +1,57 @@
+# -- coding: utf-8
+
+class KyotoTycoon
+  module Tsvrpc
+    class Skinny
+      def initialize(host, port)
+        @host = host
+        @port = port
+        @tpl = ""
+        @tpl << "POST %s HTTP/1.1\r\n"
+        @tpl << "Content-Length: %d\r\n"
+        @tpl << "Content-Type: text/tab-separated-values; colenc=%s\r\n"
+        @tpl << "\r\n%s"
+        at_exit { finish }
+      end
+
+      def request(path, params, colenc)
+        start 
+        query = KyotoTycoon::Tsvrpc.build_query(params, colenc)
+        request = @tpl % [path, query.bytesize, colenc, query]
+        @sock.write(request)
+        first_line = @sock.gets
+        status = first_line[9, 3]
+        bodylen = 0
+        body = ""
+        colenc = nil
+        loop do
+          line = @sock.gets
+          if line['Content-Type'] && line['colenc=']
+            colenc = line.match(/colenc=([A-Z])/).to_a[1]
+            next
+          end
+
+          if line['Content-Length']
+            bodylen = line.match(/[0-9]+/)[0].to_i
+            next
+          end
+
+          if line == "\r\n"
+            break
+          end
+        end
+        body = @sock.read(bodylen)
+        [status.to_i, body, colenc]
+      end
+
+      def start
+        @sock = nil if @sock && @sock.closed?
+        @sock ||= ::TCPSocket.new(@host, @port)
+      end
+
+      def finish
+        @sock.close if @sock && !@sock.closed?
+      end
+    end
+  end
+end

data/lib/kyototycoon/tsvrpc.rb

@@ -0,0 +1,45 @@
+# -- coding: utf-8
+
+
+class KyotoTycoon
+  module Tsvrpc
+    def self.parse(body, colenc)
+      decoder = case colenc
+        when "U"
+          lambda{|body| CGI.unescape(body)}
+        when "B"
+          lambda{|body| Base64.decode64(body)}
+        when nil
+          lambda{|body| body}
+        else
+          raise "Unknown colenc(response) '#{colenc}'"
+      end
+      body.split("\n").inject({}){|r, line|
+        k,v = *line.split("\t", 2).map{|v| decoder.call(v)}
+        r[k] = v
+        r
+      }
+    end
+
+    def self.build_query(params, colenc='U')
+      query = ""
+      if params
+        encoder = case colenc.to_s.upcase.to_sym
+          when :U
+            lambda{|body| CGI.escape(body.to_s)}
+          when :B
+            lambda{|body| [body.to_s].pack('m').gsub("\n","")}
+          else
+            raise "Unknown colenc '#{colenc}'"
+        end
+        query = params.inject([]){|r, tmp|
+          unless tmp.last.nil?
+            r << tmp.map{|v| encoder.call(v)}.join("\t")
+          end
+          r
+        }.join("\r\n")
+      end
+      query
+    end
+  end
+end

data/lib/kyototycoon.rb

@@ -0,0 +1,305 @@
+# -- coding: utf-8
+
+require "logger"
+require "cgi"
+require "socket"
+require "base64"
+require "timeout"
+require "kyototycoon/cursor.rb"
+require "kyototycoon/serializer.rb"
+require "kyototycoon/serializer/default.rb"
+require "kyototycoon/serializer/msgpack.rb"
+require "kyototycoon/tsvrpc.rb"
+require "kyototycoon/tsvrpc/skinny.rb"
+require "kyototycoon/stream.rb"
+
+class KyotoTycoon
+  VERSION = '0.6.0'
+
+  attr_accessor :colenc, :connect_timeout, :servers
+  attr_reader :serializer, :logger, :db
+
+  DEFAULT_HOST = '0.0.0.0'
+  DEFAULT_PORT = 1978
+
+  def self.configure(name, host=DEFAULT_HOST, port=DEFAULT_PORT, &block)
+    @configure ||= {}
+    if @configure[name]
+      raise "'#{name}' is registered"
+    end
+    @configure[name] = lambda{
+      kt = KyotoTycoon.new(host, port)
+      block.call(kt)
+      kt
+    }
+  end
+
+  def self.configures
+    @configure
+  end
+
+  def self.configure_reset!
+    @configure = {}
+  end
+  
+  def self.create(name)
+    if @configure[name].nil?
+      raise "undefined configure: '#{name}'"
+    end
+    @configure[name].call
+  end
+
+  def initialize(host=DEFAULT_HOST, port=DEFAULT_PORT)
+    @servers = [[host, port]]
+    @checked_servers = nil
+    @serializer = KyotoTycoon::Serializer::Default
+    @logger = Logger.new(nil)
+    @colenc = :B
+    @connect_timeout = 0.5
+    @cursor = 1
+  end
+
+  def serializer= (adaptor=:default)
+    klass = KyotoTycoon::Serializer.get(adaptor)
+    @serializer = klass
+  end
+
+  def db= (db)
+    @db = db
+  end
+
+  def logger= (logger)
+    if logger.class != Logger
+      logger = Logger.new(logger)
+    end
+    @logger = logger
+  end
+
+  def get(key)
+    res = request('/rpc/get', {:key => key})
+    @serializer.decode(Tsvrpc.parse(res[:body], res[:colenc])['value'])
+  end
+  alias_method :[], :get
+
+  def remove(*keys)
+    remove_bulk(keys.flatten)
+  end
+  alias_method :delete, :remove
+
+  def set(key, value, xt=nil)
+    res = request('/rpc/set', {:key => key, :value => @serializer.encode(value), :xt => xt})
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+  alias_method :[]=, :set
+
+  def add(key, value, xt=nil)
+    res = request('/rpc/add', {:key => key, :value => @serializer.encode(value), :xt => xt})
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def replace(key, value, xt=nil)
+    res = request('/rpc/replace', {:key => key, :value => @serializer.encode(value), :xt => xt})
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def append(key, value, xt=nil)
+    request('/rpc/append', {:key => key, :value => @serializer.encode(value), :xt => xt})
+  end
+
+  def cas(key, oldval, newval, xt=nil)
+    res = request('/rpc/cas', {:key => key, :oval=> @serializer.encode(oldval), :nval => @serializer.encode(newval), :xt => xt})
+    case res[:status].to_i
+      when 200
+        true
+      when 450
+        false
+    end
+  end
+
+  def increment(key, num=1, xt=nil)
+    res = request('/rpc/increment', {:key => key, :num => num, :xt => xt})
+    Tsvrpc.parse(res[:body], res[:colenc])['num'].to_i
+  end
+  alias_method :incr, :increment
+
+  def decrement(key, num=1, xt=nil)
+    increment(key, num * -1, xt)
+  end
+  alias_method :decr, :decrement
+
+  def increment_double(key, num, xt=nil)
+    res = request('/rpc/increment_double', {:key => key, :num => num, :xt => xt})
+    Tsvrpc.parse(res[:body], res[:colenc])['num'].to_f
+  end
+
+  def set_bulk(records)
+    # records={'a' => 'aa', 'b' => 'bb'}
+    values = {}
+    records.each{|k,v|
+      values[k.to_s.match(/^_/) ? k.to_s : "_#{k}"] = @serializer.encode(v)
+    }
+    res = request('/rpc/set_bulk', values)
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def get_bulk(keys)
+    params = keys.inject({}){|params, k|
+      params[k.to_s.match(/^_/) ? k.to_s : "_#{k}"] = ''
+      params
+    }
+    res = request('/rpc/get_bulk', params)
+    bulk = Tsvrpc.parse(res[:body], res[:colenc])
+    bulk.delete_if{|k,v| k.match(/^[^_]/)}.inject({}){|r, (k,v)|
+      r[k[1..-1]] = @serializer.decode(v)
+      r
+    }
+  end
+
+  def remove_bulk(keys)
+    params = keys.inject({}){|params, k|
+      params[k.to_s.match(/^_/) ? k.to_s : "_#{k}"] = ''
+      params
+    }
+    res = request('/rpc/remove_bulk', params)
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def cursor(cur_id=nil)
+    Cursor.new(self, cur_id || @cursor += 1)
+  end
+
+  def clear
+    request('/rpc/clear')
+  end
+
+  def vacuum
+    request('/rpc/vacuum')
+  end
+
+  def sync(params={})
+    request('/rpc/synchronize', params)
+  end
+  alias_method :syncronize, :sync
+
+  def echo(value)
+    res = request('/rpc/echo', value)
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def report
+    res = request('/rpc/report')
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def status
+    res = request('/rpc/status')
+    Tsvrpc.parse(res[:body], res[:colenc])
+  end
+
+  def match_prefix(prefix)
+    res = request('/rpc/match_prefix', {:prefix => prefix})
+    keys = []
+    Tsvrpc.parse(res[:body], res[:colenc]).each{|k,v|
+      if k != 'num'
+        keys << k[1, k.length]
+      end
+    }
+    keys
+  end
+
+  def match_regex(re)
+    if re.class == Regexp
+      re = re.source
+    end
+    res = request('/rpc/match_regex', {:regex => re})
+    keys = []
+    Tsvrpc.parse(res[:body], res[:colenc]).each{|k,v|
+      if k != 'num'
+        keys << k[1, k.length]
+      end
+    }
+    keys
+  end
+
+  def keys
+    match_prefix("")
+  end
+
+  def request(path, params=nil)
+    if @db
+      params ||= {}
+      params[:DB] = @db
+    end
+
+    status,body,colenc = client.request(path, params, @colenc)
+    if ![200, 450].include?(status.to_i)
+      raise body
+    end
+    res = {:status => status, :body => body, :colenc => colenc}
+    @logger.info("#{path}: #{res[:status]} with query parameters #{params.inspect}")
+    res
+  end
+
+  def client
+    host, port = *choice_server
+    @client ||= begin
+      Tsvrpc::Skinny.new(host, port)
+    end
+  end
+
+  def start
+    client.start
+  end
+
+  def finish
+    client.finish
+  end
+
+  private
+
+  def ping(host, port)
+    begin
+      rpc = Tsvrpc::Skinny.new(host, port)
+      timeout(@connect_timeout){
+        @logger.debug("connect check #{host}:#{port}")
+        res = rpc.request('/rpc/echo', {'0' => '0'}, :U)
+        @logger.debug(res)
+      }
+      true
+    rescue Timeout::Error => ex
+      # Ruby 1.8.7 compatible
+      @logger.warn("connect failed at #{host}:#{port}")
+      false
+    rescue SystemCallError
+      @logger.warn("connect failed at #{host}:#{port}")
+      false
+    rescue => ex
+      @logger.warn("connect failed at #{host}:#{port}")
+      false
+    ensure
+      # for 1.8.7
+      rpc.finish
+    end
+  end
+
+  def choice_server
+    if @checked_servers
+      return @checked_servers
+    end
+
+    @servers.each{|s|
+      host,port = *s
+      if ping(host, port)
+        @checked_servers = [host, port]
+        break
+      end
+    }
+    if @checked_servers.nil?
+      msg = "alived server not exists"
+      @logger.fatal(msg)
+      raise msg
+    end
+    @checked_servers
+  end
+
+end

data/spec/connect_spec.rb

@@ -0,0 +1,72 @@
+# -- coding: utf-8
+
+
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper.rb')
+
+describe KyotoTycoon do
+  it 'should handle multi servers' do
+    kt = KyotoTycoon.new('www.example.com', 11111)
+    kt.connect_timeout = 0.1
+    kt.servers << ['example.net', 1978]
+    kt.servers << ['0.0.0.0', 19999]
+    kt['foo'] = 'bar'
+    kt[:foo].should == 'bar'
+  end
+
+  it 'should configure/create method works' do
+    logger = Logger.new(STDOUT)
+    KyotoTycoon.configure(:test) do |kt|
+      kt.logger = logger
+      kt.serializer = :msgpack
+      kt.db = 'foobar'
+    end
+    KyotoTycoon.configure(:test2, 'host', 1999) do |kt|
+      kt.logger = logger
+      kt.serializer = :msgpack
+      kt.db = 'foobar'
+    end
+    %w!test test2!.each{|name|
+      kt = KyotoTycoon.create(name.to_sym)
+      kt.logger.should == logger
+      kt.serializer.should == KyotoTycoon::Serializer::Msgpack
+      kt.db.should == 'foobar'
+    }
+    lambda { KyotoTycoon.configure(:test2) }.should raise_error(StandardError)
+    lambda { KyotoTycoon.create(:not_exists) }.should raise_error(StandardError)
+
+    KyotoTycoon.configures.length.should == 2
+    KyotoTycoon.configure_reset!
+    KyotoTycoon.configures.length.should == 0
+  end
+
+  it 'should handle `ktremotemgr slave`' do
+    io = File.open("#{File.dirname(__FILE__)}/ktslave.txt", "r")
+    current = 0
+    KyotoTycoon::Stream.run(io){|line|
+      case current
+        when 0 # clear command
+          line.cmd.should == 'clear'
+          line.xt_time.should == Time.at(0)
+          line.value.should be_nil
+          line.key.should be_nil
+          line.value.should be_nil
+        when 1 # set foo bar
+          line.cmd.should == 'set'
+          line.xt_time.should > Time.now
+          line.key.should == 'foo'
+          line.value.should == 'bar'
+        when 2 # set fooxt bar with xt(2010-12-23 22:09:49 +0900)
+          line.cmd.should == 'set'
+          line.key.should == 'fooxt'
+          line.value.should == 'bar'
+          line.xt_time.should > Time.at(1234567890)
+          line.xt_time.should < Time.at(1334567890)
+        when 3 # remove foo
+          line.cmd.should == 'remove'
+          line.key.should == 'foo'
+          line.value.should be_nil
+      end
+      current += 1
+    }
+  end
+end

data/spec/cursor_spec.rb

@@ -0,0 +1,89 @@
+# -- coding: utf-8
+
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper.rb')
+
+describe KyotoTycoon do
+  before(:each) do
+    @kt = KyotoTycoon.new('localhost', 19999)
+    100.times{|n|
+      @kt["#{"%02d" % n}foo"] = "foo#{n}"
+    }
+  end
+
+  it 'should handle cursor object' do
+    cur = @kt.cursor
+    cur.jump("33")
+    cur.key.should == "33foo"
+    cur.value.should == "foo33"
+    cur.current.should == [ "33foo","foo33" ]
+    cur.value = "new"
+    @kt["33foo"].should == "new"
+
+    cur.step
+    cur.key.should == "34foo"
+    cur.value.should == "foo34"
+    cur.remove
+    @kt["34foo"].should be_nil
+
+    cur.jump('55')
+    cur.seize.should == ["55foo","foo55"]
+    @kt["55foo"].should be_nil
+  end
+
+  it 'should handle cursor steps' do
+    # If you got failed in this section, it's a KyotoCabinet's bug
+    # It has been fixed at KyotoCabinet 1.2.70
+    # c.f.
+    # https://gist.github.com/1117611
+    # https://twitter.com/#!/fallabs/status/98079688550916097
+    cur = @kt.cursor
+    cur.jump
+    cur.key.should == "00foo"
+
+    cur.jump_back
+    cur.key.should == "99foo"
+    cur.jump
+    cur.key.should == "00foo"
+
+    cur.jump("50")
+    cur.step
+    cur.key.should == "51foo"
+    cur.step_back
+    cur.key.should == "50foo"
+    cur.step
+    cur.key.should == "51foo"
+  end
+
+  it "should handle multiple cursors" do
+    cur = @kt.cursor
+    cur2 = @kt.cursor
+    cur.jump("33")
+    cur.cur.should_not == cur2.cur
+    cur.key.should_not == cur2.key
+  end
+
+  it "should keep current position after called #each" do
+    cur = @kt.cursor
+    cur.jump("49")
+    cur.each{|k,v| } # no-op
+    cur.key.should == "49foo"
+  end
+
+  it "should handle #each" do
+    cur = @kt.cursor
+    cur.find{|k,v| k == "non-exists"}.should be_nil
+    cur.find_all{|k,v| k.match(/3[234]foo/)}.should == [
+      %w!32foo foo32!,
+      %w!33foo foo33!,
+      %w!34foo foo34!,
+    ]
+    cur.jump("33")
+    cur.find_all{|k,v| k.match(/3[234]foo/)}.should == [
+      %w!33foo foo33!,
+      %w!34foo foo34!,
+    ]
+
+    cur.jump
+    cur.find{|k,v| k == "15foo"}.should == ["15foo", "foo15"]
+  end
+end

data/spec/ktslave.txt

@@ -0,0 +1,4 @@
+1293108572805000000	1	0	clear
+1293108575130000000	1	0	set	Zm9v	//////9iYXI=
+1293108589068000000	1	0	set	Zm9veHQ=	AE0TSh1iYXI=
+1293108597399000000	1	0	remove	Zm9v