better-riak-client 1.0.5

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

@@ -0,0 +1,111 @@
+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, options={})
+          pbuf = RpbPutReq.new(options.merge(:bucket => maybe_encode(robject.bucket.name)))
+          pbuf.key = maybe_encode(robject.key) if robject.key # Put w/o key supported!
+          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,
+                                        :indexes => robject.indexes.map {|k,s| encode_index(k,s) }.flatten)
+
+          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)
+          return robject if pbuf.respond_to?(:unchanged) && pbuf.unchanged # Reloading
+          robject.vclock = Base64.encode64(pbuf.vclock).chomp if pbuf.vclock
+          robject.key = maybe_unescape(pbuf.key) if pbuf.respond_to?(:key) && pbuf.key # Put w/o key
+          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.indexes.present?
+            robject.indexes.clear
+            pbuf.indexes.each {|pair| decode_index(pair, robject.indexes) }
+          end
+          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 decode_index(pbuf, hash)
+          value = pbuf.key =~ /int$/ ? pbuf.value.to_i : pbuf.value
+          hash[pbuf.key] << value
+        end
+
+        def encode_index(key, set)
+          set.map do |v|
+            RpbPair.new(:key => maybe_encode(key),
+                        :value => maybe_encode(v.to_s))
+          end
+        end
+
+        def maybe_encode(string)
+          ENCODING ? string.dup.force_encoding('BINARY') : string
+        end
+      end
+
+      include ObjectMethods
+    end
+  end
+end

data/lib/riak/client/beefcake_protobuffs_backend.rb

@@ -0,0 +1,226 @@
+require 'base64'
+require 'riak/json'
+require 'riak/client'
+require 'riak/failed_request'
+require 'riak/client/protobuffs_backend'
+
+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, options={})
+        options = prune_unsupported_options(:GetReq, normalize_quorums(options))
+        bucket = Bucket === bucket ? bucket.name : bucket
+        req = RpbGetReq.new(options.merge(:bucket => maybe_encode(bucket), :key => maybe_encode(key)))
+        write_protobuff(:GetReq, req)
+        decode_response(RObject.new(client.bucket(bucket), key))
+      end
+
+      def reload_object(robject, options={})
+        options = normalize_quorums(options)
+        options[:bucket] = maybe_encode(robject.bucket.name)
+        options[:key] = maybe_encode(robject.key)
+        options[:if_modified] = maybe_encode Base64.decode64(robject.vclock) if robject.vclock
+        req = RpbGetReq.new(prune_unsupported_options(:GetReq, options))
+        write_protobuff(:GetReq, req)
+        decode_response(robject)
+      end
+
+      def store_object(robject, options={})
+        options = normalize_quorums(options)
+        if robject.prevent_stale_writes
+          unless pb_conditionals?
+            other = fetch_object(robject.bucket, robject.key)
+            raise Riak::ProtobuffsFailedRequest.new(:stale_object, t("stale_write_prevented")) unless other.vclock == robject.vclock
+          end
+          if robject.vclock
+            options[:if_not_modified] = true
+          else
+            options[:if_none_match] = true
+          end
+        end
+        req = dump_object(robject, prune_unsupported_options(:PutReq, options))
+        write_protobuff(:PutReq, req)
+        decode_response(robject)
+      end
+
+      def delete_object(bucket, key, options={})
+        bucket = Bucket === bucket ? bucket.name : bucket
+        options = normalize_quorums(options)
+        options[:bucket] = maybe_encode(bucket)
+        options[:key] = maybe_encode(key)
+        options[:vclock] = Base64.decode64(options[:vclock]) if options[:vclock]
+        req = RpbDelReq.new(prune_unsupported_options(:DelReq, options))
+        write_protobuff(:DelReq, req)
+        decode_response
+      end
+
+      def get_bucket_props(bucket)
+        bucket = bucket.name if Bucket === bucket
+        req = RpbGetBucketReq.new(:bucket => maybe_encode(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 => maybe_encode(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 => maybe_encode(bucket))
+        write_protobuff(:ListKeysReq, req)
+        keys = []
+        while msg = decode_response
+          break if msg.done
+          if block_given?
+            yield msg.keys
+          else
+            keys += msg.keys
+          end
+        end
+        block_given? || keys
+      end
+
+      def mapred(mr, &block)
+        raise MapReduceError.new(t("empty_map_reduce_query")) if mr.query.empty? && !mapred_phaseless?
+        req = RpbMapRedReq.new(:request => mr.to_json, :content_type => "application/json")
+        write_protobuff(:MapRedReq, req)
+        results = []
+        while msg = decode_response
+          break if msg.done
+          if block_given?
+            yield msg.phase, JSON.parse(msg.response)
+          else
+            results[msg.phase] ||= []
+            results[msg.phase] += JSON.parse(msg.response)
+          end
+        end
+        block_given? || results.compact.size == 1 ? results.last : results
+      end
+
+      def get_index(bucket, index, query)
+        return super unless pb_indexes?
+        if Range === query
+          options = {
+            :qtype => RpbIndexReq::IndexQueryType::RANGE,
+            :range_min => query.begin.to_s,
+            :range_max => query.end.to_s
+          }
+        else
+          options = {
+            :qtype => RpbIndexReq::IndexQueryType::EQ,
+            :key => query.to_s
+          }
+        end
+        req = RpbIndexReq.new(options.merge(:bucket => bucket, :index => index))
+        write_protobuff(:IndexReq, req)
+        decode_response
+      end
+
+      def search(index, query, options={})
+        return super unless pb_search?
+        options = options.symbolize_keys
+        options[:op] = options.delete(:'q.op') if options[:'q.op']
+        req = RpbSearchQueryReq.new(options.merge(:index => index || 'search', :q => query))
+        write_protobuff(:SearchQueryReq, req)
+        decode_response
+      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
+            res = RpbGetResp.decode(message)
+            load_object(res, args.first)
+          when :PutResp
+            res = RpbPutResp.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)
+          when :IndexResp
+            res = RpbIndexResp.decode(message)
+            res.keys
+          when :SearchQueryResp
+            res = RpbSearchQueryResp.decode(message)
+            { 'docs' => res.docs.map {|d| decode_doc(d) },
+              'max_score' => res.max_score,
+              'num_found' => res.num_found }
+          end
+        end
+      rescue SystemCallError, SocketError => e
+        reset_socket
+        raise
+      end
+
+      def decode_doc(doc)
+        Hash[doc.properties.map {|p| [ p.key, p.value ] }]
+      end
+    end
+  end
+end

data/lib/riak/client/decaying.rb

@@ -0,0 +1,36 @@
+module Riak
+  class Client
+    # A float value which decays exponentially toward 0 over time.
+    # @private
+    class Decaying
+      attr_accessor :e
+      attr_accessor :p
+
+      # @param [Hash] opts options
+      # @option options [Float] :p (0.0) The initial value
+      # @option options [Float] :e (Math::E) Exponent base
+      # @option options [Float] :r (Math.log(0.5) / 10) Timescale
+      #   factor - defaulting to decay 50% every 10 seconds
+      def initialize(opts = {})
+        @p = opts[:p] || 0.0
+        @e = opts[:e] || Math::E
+        @r = opts[:r] || Math.log(0.5) / 10
+        @t0 = Time.now
+      end
+
+      # Add to current value.
+      # @param [Float] d the value to add
+      def <<(d)
+        @p = value + d
+      end
+
+      # @return [Float] the current value (adjusted for the time decay)
+      def value
+        now = Time.now
+        dt = now - @t0
+        @t0 = now
+        @p = @p * (@e ** (@r * dt))
+      end
+    end
+  end
+end

data/lib/riak/client/excon_backend.rb

@@ -0,0 +1,162 @@
+require 'riak/failed_request'
+require 'riak/client/http_backend'
+require 'riak/client/http_backend/request_headers'
+
+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'
+          minimum_version?("0.5.7") && register_exceptions && handle_deprecations && patch_sockets
+        rescue LoadError
+          false
+        end
+      end
+
+      # Adds Excon's relevant internal exceptions to the rescuable
+      # network-related errors.
+      def self.register_exceptions
+        unless Client::NETWORK_ERRORS.include?(Excon::Errors::SocketError)
+          Client::NETWORK_ERRORS << Excon::Errors::SocketError
+          Client::NETWORK_ERRORS << Excon::Errors::Timeout if defined? Excon::Errors::Timeout
+        end
+        true
+      end
+
+      # Adjusts Excon's connection collection to allow multiple
+      # connections to the same host from the same Thread. Instead we
+      # use the Riak::Client::Pool to segregate connections.
+      # @note This can be changed when Excon has a proper pool of its own.
+      def self.patch_sockets
+        unless defined? @@patched
+          ::Excon::Connection.class_eval do
+            def sockets
+              @sockets ||= {}
+            end
+          end
+        end
+        @@patched = true
+      end
+
+      # Defines instance methods that handle changes in the Excon API
+      # across different versions.
+      def self.handle_deprecations
+        # Define #make_request
+        unless method_defined?(:make_request)
+          if minimum_version?("0.10.2")
+            def make_request(params, block)
+              params[:response_block] = block if block
+              connection.request(params)
+            end
+          else
+            def make_request(params, block)
+              response = connection.request(params, &block)
+            end
+          end
+          private :make_request
+        end
+
+        # Define #configure_ssl
+        unless method_defined?(:configure_ssl)
+          if minimum_version?("0.9.6")
+            def configure_ssl
+              Excon.defaults[:ssl_verify_peer] = (@node.ssl_options[:verify_mode].to_s === "peer")
+              Excon.defaults[:ssl_ca_path]     = @node.ssl_options[:ca_path] if @node.ssl_options[:ca_path]
+            end
+          else
+            def configure_ssl
+              Excon.ssl_verify_peer = (@node.ssl_options[:verify_mode].to_s === "peer")
+              Excon.ssl_ca_path     = @node.ssl_options[:ca_path] if @node.ssl_options[:ca_path]
+            end
+          end
+          private :configure_ssl
+        end
+        true
+      end
+
+      # Returns true if the Excon library is at least the given
+      # version. This is used inside the backend to check how to
+      # provide certain request and configuration options.
+      def self.minimum_version?(version)
+        Gem::Version.new(Excon::VERSION) >= Gem::Version.new(version)
+      end
+
+      # Sets the connect timeout applied to the Excon connection
+      # Increase this if you have very long request times.
+      def self.connect_timeout=(timeout)
+        @connect_timeout = timeout
+      end
+
+      def self.connect_timeout
+        @connect_timeout ||= 4096
+      end
+
+      # Sets the read_timeout applied to the Excon connection
+      # Increase this if you have very long request times.
+      def self.read_timeout=(timeout)
+        @read_timeout = timeout
+      end
+
+      def self.read_timeout
+        @read_timeout ||= 4096
+      end
+
+      # Sets the write_timeout applied to the Excon connection
+      # Increase this if you have very long request times.
+      def self.write_timeout=(timeout)
+        @write_timeout = timeout
+      end
+
+      def self.write_timeout
+        @write_timeout ||= 4096
+      end
+
+      def teardown
+        connection.reset
+      end
+
+      private
+      def perform(method, uri, headers, expect, data=nil, &block)
+        configure_ssl if @node.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)
+
+        # Later versions of Excon pass multiple arguments to the block
+        block = lambda {|*args| yield args.first } if block_given?
+
+        response = make_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,
+          :read_timeout => self.class.read_timeout,
+          :write_timeout => self.class.write_timeout,
+          :connect_timeout => self.class.connect_timeout
+        )
+      end
+    end
+  end
+end