seomoz-riak-client 1.0.0.pre

data/lib/riak/client/beefcake/messages.rb

@@ -0,0 +1,147 @@
+require 'beefcake'
+
+module Riak
+  class Client
+    # @private
+    class BeefcakeProtobuffsBackend
+      # Embedded messages
+      class RpbPair
+        include Beefcake::Message
+        required :key,   :bytes, 1
+        optional :value, :bytes, 2
+      end
+
+      class RpbBucketProps
+        include Beefcake::Message
+        optional :n_val,      :uint32, 1
+        optional :allow_mult, :bool,   2
+      end
+
+      class RpbLink
+        include Beefcake::Message
+        optional :bucket, :bytes, 1
+        optional :key,    :bytes, 2
+        optional :tag,    :bytes, 3
+      end
+
+      class RpbContent
+        include Beefcake::Message
+        required :value,            :bytes,  1
+        optional :content_type,     :bytes,  2
+        optional :charset,          :bytes,  3
+        optional :content_encoding, :bytes,  4
+        optional :vtag,             :bytes,  5
+        repeated :links,            RpbLink, 6
+        optional :last_mod,         :uint32, 7
+        optional :last_mod_usecs,   :uint32, 8
+        repeated :usermeta,         RpbPair, 9
+      end
+
+      # Primary messages
+      class RpbErrorResp
+        include Beefcake::Message
+        required :errmsg,  :bytes,  1
+        required :errcode, :uint32, 2
+      end
+
+      class RpbGetClientIdResp
+        include Beefcake::Message
+        required :client_id, :bytes, 1
+      end
+
+      class RpbSetClientIdReq
+        include Beefcake::Message
+        required :client_id, :bytes, 1
+      end
+
+      class RpbGetServerInfoResp
+        include Beefcake::Message
+        optional :node,           :bytes, 1
+        optional :server_version, :bytes, 2
+      end
+
+      class RpbGetReq
+        include Beefcake::Message
+        required :bucket, :bytes,  1
+        required :key,    :bytes,  2
+        optional :r,      :uint32, 3
+      end
+
+      class RpbGetResp
+        include Beefcake::Message
+        repeated :content, RpbContent, 1
+        optional :vclock,  :bytes,  2
+      end
+
+      class RpbPutReq
+        include Beefcake::Message
+        required :bucket,      :bytes,  1
+        required :key,         :bytes,  2
+        optional :vclock,      :bytes,  3
+        required :content,     RpbContent, 4
+        optional :w,           :uint32, 5
+        optional :dw,          :uint32, 6
+        optional :return_body, :bool,   7
+      end
+
+      # Optional since it has the same structure as GetResp
+      # class RpbPutResp
+      #   include Beefcake::Message
+      #   repeated :content, RpbContent, 1
+      #   optional :vclock,  :bytes,  2
+      # end
+
+      class RpbDelReq
+        include Beefcake::Message
+        required :bucket, :bytes,  1
+        required :key,    :bytes,  2
+        optional :rw,     :uint32, 3
+      end
+
+      class RpbListBucketsResp
+        include Beefcake::Message
+        repeated :buckets, :bytes, 1
+      end
+
+      class RpbListKeysReq
+        include Beefcake::Message
+        required :bucket, :bytes, 1
+      end
+
+      class RpbListKeysResp
+        include Beefcake::Message
+        repeated :keys, :bytes, 1
+        optional :done, :bool,  2
+      end
+
+      class RpbGetBucketReq
+        include Beefcake::Message
+        required :bucket, :bytes, 1
+      end
+
+      class RpbGetBucketResp
+        include Beefcake::Message
+        required :props, RpbBucketProps, 1
+      end
+
+      class RpbSetBucketReq
+        include Beefcake::Message
+        required :bucket, :bytes,         1
+        required :props,  RpbBucketProps, 2
+      end
+
+      class RpbMapRedReq
+        include Beefcake::Message
+        required :request,      :bytes, 1
+        required :content_type, :bytes, 2
+      end
+
+      class RpbMapRedResp
+        include Beefcake::Message
+        optional :phase,    :uint32, 1
+        optional :response, :bytes,  2
+        optional :done,     :bool,   3
+      end
+    end
+  end
+end

data/lib/riak/client/beefcake/object_methods.rb

@@ -0,0 +1,92 @@
+require 'riak/robject'
+require 'riak/link'
+require 'riak/client/beefcake/messages'
+
+module Riak
+  class Client
+    class BeefcakeProtobuffsBackend
+      module ObjectMethods
+        ENCODING = "Riak".respond_to?(:encoding)
+
+        # Returns RpbPutReq
+        def dump_object(robject)
+          pbuf = RpbPutReq.new(:bucket => maybe_encode(robject.bucket.name))
+          pbuf.key = maybe_encode(robject.key ||= generate_key)
+          pbuf.vclock = maybe_encode Base64.decode64(robject.vclock) if robject.vclock
+          pbuf.content = RpbContent.new(:value => maybe_encode(robject.raw_data),
+                                        :content_type => maybe_encode(robject.content_type),
+                                        :links => robject.links.map {|l| encode_link(l) }.compact)
+
+          pbuf.content.usermeta = robject.meta.map {|k,v| encode_meta(k,v)} if robject.meta.any?
+          pbuf.content.vtag = maybe_encode(robject.etag) if robject.etag.present?
+          if ENCODING # 1.9 support
+            pbuf.content.charset = maybe_encode(robject.raw_data.encoding.name)
+          end
+          pbuf
+        end
+
+        # Returns RObject
+        def load_object(pbuf, robject)
+          robject.vclock = Base64.encode64(pbuf.vclock).chomp if pbuf.vclock
+          if pbuf.content.size > 1
+            robject.conflict = true
+            robject.siblings = pbuf.content.map do |c|
+              sibling = RObject.new(robject.bucket, robject.key)
+              sibling.vclock = robject.vclock
+              load_content(c, sibling)
+            end
+
+            return robject.attempt_conflict_resolution
+          else
+            load_content(pbuf.content.first, robject)
+          end
+          robject
+        end
+
+        private
+        def load_content(pbuf, robject)
+          if ENCODING && pbuf.charset.present?
+            pbuf.value.force_encoding(pbuf.charset) if Encoding.find(pbuf.charset)
+          end
+          robject.raw_data = pbuf.value
+          robject.etag = pbuf.vtag if pbuf.vtag.present?
+          robject.content_type = pbuf.content_type if pbuf.content_type.present?
+          robject.links = pbuf.links.map(&method(:decode_link)) if pbuf.links.present?
+          pbuf.usermeta.each {|pair| decode_meta(pair, robject.meta) } if pbuf.usermeta.present?
+          if pbuf.last_mod.present?
+            robject.last_modified = Time.at(pbuf.last_mod)
+            robject.last_modified += pbuf.last_mod_usecs / 1000000 if pbuf.last_mod_usecs.present?
+          end
+          robject
+        end
+
+        def decode_link(pbuf)
+          Riak::Link.new(pbuf.bucket, pbuf.key, pbuf.tag)
+        end
+
+        def encode_link(link)
+          return nil unless link.key.present?
+          RpbLink.new(:bucket => maybe_encode(link.bucket.to_s),
+                      :key => maybe_encode(link.key.to_s),
+                      :tag => maybe_encode(link.tag.to_s))
+        end
+
+        def decode_meta(pbuf, hash)
+          hash[pbuf.key] = pbuf.value
+        end
+
+        def encode_meta(key,value)
+          return nil unless value.present?
+          RpbPair.new(:key => maybe_encode(key.to_s),
+                      :value => maybe_encode(value.to_s))
+        end
+
+        def maybe_encode(string)
+          ENCODING ? string.encode('BINARY') : string
+        end
+      end
+
+      include ObjectMethods
+    end
+  end
+end

data/lib/riak/client/beefcake_protobuffs_backend.rb

@@ -0,0 +1,176 @@
+
+require 'riak/json'
+require 'riak/client'
+require 'riak/failed_request'
+require 'riak/client/protobuffs_backend'
+require 'riak/client/pump'
+
+module Riak
+  class Client
+    class BeefcakeProtobuffsBackend < ProtobuffsBackend
+      def self.configured?
+        begin
+          require 'beefcake'
+          require 'riak/client/beefcake/messages'
+          require 'riak/client/beefcake/object_methods'
+          true
+        rescue LoadError, NameError
+          false
+        end
+      end
+
+      def set_client_id(id)
+        value = case id
+                when Integer
+                  [id].pack("N")
+                else
+                  id.to_s
+                end
+        req = RpbSetClientIdReq.new(:client_id => value)
+        write_protobuff(:SetClientIdReq, req)
+        decode_response
+      end
+
+      def fetch_object(bucket, key, r=nil)
+        bucket = Bucket === bucket ? bucket.name : bucket
+        req = RpbGetReq.new(:bucket => bucket, :key => key)
+        req.r = normalize_quorum_value(r) if r
+        write_protobuff(:GetReq, req)
+        decode_response(RObject.new(client.bucket(bucket), key))
+      end
+
+      def reload_object(robject, r=nil)
+        req = RpbGetReq.new(:bucket => robject.bucket.name, :key => robject.key)
+        req.r = normalize_quorum_value(r) if r
+        write_protobuff(:GetReq, req)
+        decode_response(robject)
+      end
+
+      def store_object(robject, returnbody=false, w=nil, dw=nil)
+        if robject.prevent_stale_writes
+          other = fetch_object(robject.bucket, robject.key)
+          raise Riak::ProtobuffsFailedRequest(:stale_object, t("stale_write_prevented")) unless other.vclock == robject.vclock
+        end
+        req = dump_object(robject)
+        req.w = normalize_quorum_value(w) if w
+        req.dw = normalize_quorum_value(dw) if dw
+        req.return_body = returnbody
+        write_protobuff(:PutReq, req)
+        decode_response(robject)
+      end
+
+      def delete_object(bucket, key, rw=nil)
+        bucket = Bucket === bucket ? bucket.name : bucket
+        req = RpbDelReq.new(:bucket => bucket, :key => key)
+        req.rw = normalize_quorum_value(rw) if rw
+        write_protobuff(:DelReq, req)
+        decode_response
+      end
+
+      def get_bucket_props(bucket)
+        bucket = bucket.name if Bucket === bucket
+        req = RpbGetBucketReq.new(:bucket => bucket)
+        write_protobuff(:GetBucketReq, req)
+        decode_response
+      end
+
+      def set_bucket_props(bucket, props)
+        bucket = bucket.name if Bucket === bucket
+        props = props.slice('n_val', 'allow_mult')
+        req = RpbSetBucketReq.new(:bucket => bucket, :props => RpbBucketProps.new(props))
+        write_protobuff(:SetBucketReq, req)
+        decode_response
+      end
+
+      def list_keys(bucket, &block)
+        bucket = bucket.name if Bucket === bucket
+        req = RpbListKeysReq.new(:bucket => bucket)
+        write_protobuff(:ListKeysReq, req)
+        keys = []
+        pump = Pump.new(block) if block_given?
+        while msg = decode_response
+          break if msg.done
+          if pump
+            pump.pump msg.keys
+          else
+            keys += msg.keys
+          end
+        end
+        block_given? || keys
+      end
+
+      def mapred(mr, &block)
+        req = RpbMapRedReq.new(:request => mr.to_json, :content_type => "application/json")
+        write_protobuff(:MapRedReq, req)
+        results = []
+        pump = Pump.new(lambda do |msg|
+          block.call msg.phase, JSON.parse(msg.response)
+        end) if block_given?
+        while msg = decode_response
+          break if msg.done
+          if pump
+            pump.pump msg
+          else
+            results[msg.phase] ||= []
+            results[msg.phase] += JSON.parse(msg.response)
+          end
+        end
+        block_given? || results.compact.size == 1 ? results.last : results
+      end
+
+      private
+      def write_protobuff(code, message)
+        encoded = message.encode
+        header = [encoded.length+1, MESSAGE_CODES.index(code)].pack("NC")
+        socket.write(header + encoded)
+      end
+
+      def decode_response(*args)
+        header = socket.read(5)
+        raise SocketError, "Unexpected EOF on PBC socket" if header.nil?
+        msglen, msgcode = header.unpack("NC")
+        if msglen == 1
+          case MESSAGE_CODES[msgcode]
+          when :PingResp, :SetClientIdResp, :PutResp, :DelResp, :SetBucketResp
+            true
+          when :ListBucketsResp, :ListKeysResp
+            []
+          when :GetResp
+            raise Riak::ProtobuffsFailedRequest.new(:not_found, t('not_found'))
+          else
+            false
+          end
+        else
+          message = socket.read(msglen-1)
+          case MESSAGE_CODES[msgcode]
+          when :ErrorResp
+            res = RpbErrorResp.decode(message)
+            raise Riak::ProtobuffsFailedRequest.new(res.errcode, res.errmsg)
+          when :GetClientIdResp
+            res = RpbGetClientIdResp.decode(message)
+            res.client_id
+          when :GetServerInfoResp
+            res = RpbGetServerInfoResp.decode(message)
+            {:node => res.node, :server_version => res.server_version}
+          when :GetResp, :PutResp
+            res = RpbGetResp.decode(message)
+            load_object(res, args.first)
+          when :ListBucketsResp
+            res = RpbListBucketsResp.decode(message)
+            res.buckets
+          when :ListKeysResp
+            RpbListKeysResp.decode(message)
+          when :GetBucketResp
+            res = RpbGetBucketResp.decode(message)
+            {'n_val' => res.props.n_val, 'allow_mult' => res.props.allow_mult}
+          when :MapRedResp
+            RpbMapRedResp.decode(message)
+          end
+        end
+      rescue SocketError => e
+        reset_socket
+        raise Riak::ProtobuffsFailedRequest.new(:server_error, e.message)
+      end
+    end
+  end
+end

data/lib/riak/client/excon_backend.rb

@@ -0,0 +1,65 @@
+
+require 'riak/failed_request'
+require 'riak/client/http_backend'
+require 'riak/client/http_backend/request_headers'
+require 'riak/client/pump'
+
+module Riak
+  class Client
+    # An HTTP backend for Riak::Client that uses Wesley Beary's Excon
+    # HTTP library. Conforms to the Riak::Client::HTTPBackend
+    # interface.
+    class ExconBackend < HTTPBackend
+      def self.configured?
+        begin
+          require 'excon'
+          Excon::VERSION >= "0.5.7"
+        rescue LoadError
+          false
+        end
+      end
+
+      private
+      def perform(method, uri, headers, expect, data=nil, &block)
+        configure_ssl if @client.ssl_enabled?
+
+        params = {
+          :method => method.to_s.upcase,
+          :headers => RequestHeaders.new(headers).to_hash,
+          :path => uri.path
+        }
+        params[:query] = uri.query if uri.query
+        params[:body] = data if [:put,:post].include?(method)
+        params[:idempotent] = (method != :post)
+
+        if block_given?
+          pump = Pump.new(block)
+          # Later versions of Excon pass multiple arguments to the block
+          block = lambda {|*args| pump.pump(args.first) }
+        end
+
+        response = connection.request(params, &block)
+        response_headers.initialize_http_header(response.headers)
+
+        if valid_response?(expect, response.status)
+          result = {:headers => response_headers.to_hash, :code => response.status}
+          if return_body?(method, response.status, block_given?)
+            result[:body] = response.body
+          end
+          result
+        else
+          raise HTTPFailedRequest.new(method, expect, response.status, response_headers.to_hash, response.body)
+        end
+      end
+
+      def connection
+        @connection ||= Excon::Connection.new(root_uri.to_s)
+      end
+
+      def configure_ssl
+        Excon.ssl_verify_peer = @client.ssl_options[:verify_mode].to_s === "peer"
+        Excon.ssl_ca_path     = @client.ssl_options[:ca_path] if @client.ssl_options[:ca_path]
+      end
+    end
+  end
+end