better-riak-client 1.0.5

data/lib/riak/client/net_http_backend.rb

@@ -0,0 +1,82 @@
+
+require 'riak/client/http_backend'
+require 'riak/failed_request'
+
+module Riak
+  class Client
+    # Uses the Ruby standard library Net::HTTP to connect to Riak.
+    # Conforms to the Riak::Client::HTTPBackend interface.
+    class NetHTTPBackend < HTTPBackend
+      def self.configured?
+        begin
+          require 'net/http'
+          require 'openssl'
+          true
+        rescue LoadError, NameError
+          false
+        end
+      end
+
+      # Sets the read_timeout applied to Net::HTTP connections
+      # 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
+
+      # Net::HTTP doesn't use persistent connections, so there's no
+      # work to do here.
+      def teardown; end
+
+      private
+      def perform(method, uri, headers, expect, data=nil) #:nodoc:
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.read_timeout = self.class.read_timeout
+        configure_ssl(http) if @node.ssl_enabled?
+
+        request = Net::HTTP.const_get(method.to_s.capitalize).new(uri.request_uri, headers)
+        if String === data
+          request.body = data
+        elsif data.respond_to?(:read)
+          case
+          when data.respond_to?(:stat) # IO#stat
+            request.content_length = data.stat.size
+          when data.respond_to?(:size) # Some IO-like objects
+            request.content_length = data.size
+          else
+            request['Transfer-Encoding'] = 'chunked'
+          end
+          request.body_stream = data
+        end
+
+        {}.tap do |result|
+          http.request(request) do |response|
+            unless valid_response?(expect, response.code)
+              raise Riak::HTTPFailedRequest.new(method, expect, response.code.to_i, response.to_hash, response.body)
+            end
+
+            result.merge!({:headers => response.to_hash, :code => response.code.to_i})
+            response.read_body {|chunk| yield chunk } if block_given?
+            if return_body?(method, response.code, block_given?)
+              result[:body] = response.body
+            end
+          end
+        end
+      end
+
+      def configure_ssl(http)
+        http.use_ssl = true
+        http.verify_mode = OpenSSL::SSL.const_get("VERIFY_#{@node.ssl_options[:verify_mode].upcase}")
+        if @node.ssl_options[:pem]
+          http.cert = OpenSSL::X509::Certificate.new(@node.ssl_options[:pem])
+          http.key  = OpenSSL::PKey::RSA.new(@node.ssl_options[:pem], @node.ssl_options[:pem_password])
+        end
+        http.ca_file = @node.ssl_options[:ca_file] if @node.ssl_options[:ca_file]
+        http.ca_path = @node.ssl_options[:ca_path] if @node.ssl_options[:ca_path]
+      end
+    end
+  end
+end

data/lib/riak/client/node.rb

@@ -0,0 +1,115 @@
+module Riak
+  class Client
+    class Node
+      # Represents a single riak node in a cluster.
+
+      include Util::Translation
+      include Util::Escape
+
+      VALID_OPTIONS = [:host, :http_port, :pb_port, :http_paths, :prefix,
+        :mapred, :luwak, :solr, :port, :basic_auth, :ssl_options, :ssl]
+
+      # For a score which halves in 10 seconds, choose
+      # ln(1/2)/10
+      ERRORS_DECAY_RATE = Math.log(0.5)/10
+
+      # What IP address or hostname does this node listen on?
+      attr_accessor :host
+      # Which port does the HTTP interface listen on?
+      attr_accessor :http_port
+      # Which port does the protocol buffers interface listen on?
+      attr_accessor :pb_port
+      # A hash of HTTP paths used on this node.
+      attr_accessor :http_paths
+      # A "user:password" string.
+      attr_reader :basic_auth
+      attr_accessor :ssl_options
+      # A Decaying rate of errors.
+      attr_reader :error_rate
+
+      def initialize(client, opts = {})
+        @client = client
+        @ssl = opts[:ssl]
+        @ssl_options = opts[:ssl_options]
+        @host = opts[:host] || "127.0.0.1"
+        @http_port = opts[:http_port] || opts[:port] || 8098
+        @pb_port = opts[:pb_port] || 8087
+        @http_paths = {
+          :prefix => opts[:prefix] || "/riak/",
+          :mapred => opts[:mapred] || "/mapred",
+          :luwak =>  opts[:luwak]  || "/luwak",
+          :solr =>   opts[:solr]   || "/solr" # Unused?
+        }.merge(opts[:http_paths] || {})
+        self.basic_auth = opts[:basic_auth]
+
+        @error_rate = Decaying.new
+      end
+
+      def ==(o)
+        o.kind_of? Node and
+          @host == o.host and
+          @http_port == o.http_port and
+          @pb_port == o.pb_port
+      end
+
+      # Sets the HTTP Basic Authentication credentials.
+      # @param [String] value an auth string in the form "user:password"
+      def basic_auth=(value)
+        case value
+        when nil
+          @basic_auth = nil
+        when String
+          raise ArgumentError, t("invalid_basic_auth") unless value.to_s.split(':').length === 2
+          @basic_auth = value
+        end
+      end
+
+      # Can this node be used for HTTP requests?
+      def http?
+        # TODO: Need to sort out capabilities
+        true
+      end
+
+      # Can this node be used for protocol buffers requests?
+      def protobuffs?
+        # TODO: Need to sort out capabilities
+        true
+      end
+
+      # Enables or disables SSL on this node to be utilized by the HTTP
+      # Backends
+      def ssl=(value)
+        @ssl_options ||= Hash === value ? value : {}
+        value ? ssl_enable : ssl_disable
+      end
+
+      # Checks if SSL is enabled for HTTP
+      def ssl_enabled?
+        @client.protocol == 'https' && @ssl_options.present?
+      end
+
+      def inspect
+        "#<Node #{@host}:#{@http_port}:#{@pb_port}>"
+      end
+
+      protected
+
+
+      def ssl_enable
+        @client.protocol = 'https' unless @client.protocol == 'https'
+        @ssl_options[:pem] = File.read(@ssl_options[:pem_file]) if @ssl_options[:pem_file]
+        @ssl_options[:verify_mode] ||= "peer" if @ssl_options.stringify_keys.any? {|k,v| %w[pem ca_file ca_path].include?(k)}
+        @ssl_options[:verify_mode] ||= "none"
+        raise ArgumentError.new(t('invalid_ssl_verify_mode', :invalid => @ssl_options[:verify_mode])) unless %w[none peer].include?(@ssl_options[:verify_mode])
+
+        @ssl_options
+      end
+
+      def ssl_disable
+        @client.protocol = 'http' unless @client.protocol == 'http'
+        @ssl_options  = nil
+      end
+
+    end
+  end
+end

data/lib/riak/client/protobuffs_backend.rb

@@ -0,0 +1,173 @@
+require 'riak'
+require 'socket'
+require 'base64'
+require 'digest/sha1'
+require 'riak/util/translation'
+require 'riak/client/feature_detection'
+
+module Riak
+  class Client
+    class ProtobuffsBackend
+      include Util::Translation
+      include Util::Escape
+      include FeatureDetection
+
+      # Message Codes Enum
+      MESSAGE_CODES = %W[
+          ErrorResp
+          PingReq
+          PingResp
+          GetClientIdReq
+          GetClientIdResp
+          SetClientIdReq
+          SetClientIdResp
+          GetServerInfoReq
+          GetServerInfoResp
+          GetReq
+          GetResp
+          PutReq
+          PutResp
+          DelReq
+          DelResp
+          ListBucketsReq
+          ListBucketsResp
+          ListKeysReq
+          ListKeysResp
+          GetBucketReq
+          GetBucketResp
+          SetBucketReq
+          SetBucketResp
+          MapRedReq
+          MapRedResp
+          IndexReq
+          IndexResp
+          SearchQueryReq
+          SearchQueryResp
+       ].map {|s| s.intern }.freeze
+
+      def self.simple(method, code)
+        define_method method do
+          socket.write([1, MESSAGE_CODES.index(code)].pack('NC'))
+          decode_response
+        end
+      end
+
+      attr_accessor :client
+      attr_accessor :node
+      def initialize(client, node)
+        @client = client
+        @node = node
+      end
+
+      simple :ping,          :PingReq
+      simple :get_client_id, :GetClientIdReq
+      simple :server_info,   :GetServerInfoReq
+      simple :list_buckets,  :ListBucketsReq
+
+      # Performs a secondary-index query via emulation through MapReduce.
+      # @param [String, Bucket] bucket the bucket to query
+      # @param [String] index the index to query
+      # @param [String, Integer, Range] query the equality query or
+      #   range query to perform
+      # @return [Array<String>] a list of keys matching the query
+      def get_index(bucket, index, query)
+        mr = Riak::MapReduce.new(client).index(bucket, index, query)
+        unless mapred_phaseless?
+          mr.reduce(%w[riak_kv_mapreduce reduce_identity], :arg => {:reduce_phase_only_1 => true}, :keep => true)
+        end
+        mapred(mr).map {|p| p.last }
+      end
+
+      # Performs search query via emulation through MapReduce. This
+      # has more limited capabilites than native queries. Essentially,
+      # only the 'id' field of matched documents will ever be
+      # returned, the 'fl' and other options have no effect.
+      # @param [String] index the index to query
+      # @param [String] query the Lucene-style search query
+      # @param [Hash] options ignored in MapReduce emulation
+      # @return [Hash] the search results
+      def search(index, query, options={})
+        mr = Riak::MapReduce.new(client).search(index || 'search', query)
+        unless mapred_phaseless?
+          mr.reduce(%w[riak_kv_mapreduce reduce_identity], :arg => {:reduce_phase_only_1 => true}, :keep => true)
+        end
+        docs = mapred(mr).map {|d| {'id' => d[1] } }
+        # Since we don't get this information back from the MapReduce,
+        # we have to fake the max_score and num_found.
+        { 'docs' => docs,
+          'num_found' => docs.size,
+          'max_score' => 0.0 }
+      end
+
+      # Gracefully shuts down this connection.
+      def teardown
+        reset_socket
+      end
+
+      private
+      def get_server_version
+        server_info[:server_version]
+      end
+
+      # Implemented by subclasses
+      def decode_response
+        raise NotImplementedError
+      end
+
+      def socket
+        @socket ||= new_socket
+      end
+
+      def new_socket
+        socket = TCPSocket.new(@node.host, @node.pb_port)
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, true)
+        #TODO: Should we set the client ID here?
+        # set_client_id @client.client_id
+        socket
+      end
+
+      def reset_socket
+        reset_server_version
+        @socket.close if @socket && !@socket.closed?
+        @socket = nil
+      end
+
+      UINTMAX = 0xffffffff
+      QUORUMS = {
+        "one" => UINTMAX - 1,
+        "quorum" => UINTMAX - 2,
+        "all" => UINTMAX - 3,
+        "default" => UINTMAX - 4
+      }.freeze
+
+      def prune_unsupported_options(req,options={})
+        unless quorum_controls?
+          [:notfound_ok, :basic_quorum, :pr, :pw].each {|k| options.delete k }
+        end
+        unless pb_head?
+          [:head, :return_head].each {|k| options.delete k }
+        end
+        unless tombstone_vclocks?
+          options.delete :deletedvclock
+          options.delete :vclock if req == :DelReq
+        end
+        unless pb_conditionals?
+          [:if_not_modified, :if_none_match, :if_modified].each {|k| options.delete k }
+        end
+        options
+      end
+
+      def normalize_quorums(options={})
+        options.dup.tap do |o|
+          [:r, :pr, :w, :pw, :dw, :rw].each do |k|
+            o[k] = normalize_quorum_value(o[k]) if o[k]
+          end
+        end
+      end
+
+      def normalize_quorum_value(q)
+        QUORUMS[q.to_s] || q.to_i
+      end
+    end
+  end
+end

data/lib/riak/client/search.rb

@@ -0,0 +1,91 @@
+require 'builder'
+
+module Riak
+  class Client
+    # (Riak Search) Performs a search via the Solr interface.
+    # @overload search(index, query, options={})
+    #   @param [String] index the index to query on
+    #   @param [String] query a Lucene query string
+    # @overload search(query, options={})
+    #   Queries the default index
+    #   @param [String] query a Lucene query string
+    # @param [Hash] options extra options for the Solr query
+    # @option options [String] :df the default field to search in
+    # @option options [String] :'q.op' the default operator between terms ("or", "and")
+    # @option options [String] :wt ("json") the response type - "json" and "xml" are valid
+    # @option options [String] :sort ('none') the field and direction to sort, e.g. "name asc"
+    # @option options [Fixnum] :start (0) the offset into the query to start from, e.g. for pagination
+    # @option options [Fixnum] :rows (10) the number of results to return
+    # @return [Hash] the query result, containing the 'responseHeaders' and 'response' keys
+    def search(*args)
+      options = args.extract_options!
+      index, query = args[-2], args[-1]  # Allows nil index, while keeping it as firstargument
+      backend do |b|
+        b.search(index, query, options)
+      end
+    end
+    alias :select :search
+
+    # (Riak Search) Adds documents to a search index via the Solr interface.
+    # @overload index(index, *docs)
+    #   Adds documents to the specified search index
+    #   @param [String] index the index in which to add/update the given documents
+    #   @param [Array<Hash>] docs unnested document hashes, with one key per field
+    # @overload index(*docs)
+    #   Adds documents to the default search index
+    #   @param [Array<Hash>] docs unnested document hashes, with one key per field
+    # @raise [ArgumentError] if any documents don't include 'id' key
+    def index(*args)
+      index = args.shift if String === args.first # Documents must be hashes of fields
+      raise ArgumentError.new(t("search_docs_require_id")) unless args.all? {|d| d.key?("id") || d.key?(:id) }
+      xml = Builder::XmlMarkup.new
+      xml.add do
+        args.each do |doc|
+          xml.doc do
+            doc.each do |k,v|
+              xml.field('name' => k.to_s) { xml.text!(v.to_s) }
+            end
+          end
+        end
+      end
+      http do |h|
+        h.update_search_index(index, xml.target!)
+      end
+      true
+    end
+    alias :add_doc :index
+
+    # (Riak Search) Removes documents from a search index via the Solr interface.
+    # @overload remove(index, specs)
+    #   Removes documents from the specified index
+    #   @param [String] index the index from which to remove documents
+    #   @param [Array<Hash>] specs the specificaiton of documents to remove (must contain 'id' or 'query' keys)
+    # @overload remove(specs)
+    #   Removes documents from the default index
+    #   @param [Array<Hash>] specs the specification of documents to remove (must contain 'id' or 'query' keys)
+    # @raise [ArgumentError] if any document specs don't include 'id' or 'query' keys
+    def remove(*args)
+      index = args.shift if String === args.first
+      raise ArgumentError.new(t("search_remove_requires_id_or_query")) unless args.all? { |s|
+        s.include? :id or
+        s.include? 'id' or
+        s.include? :query or
+        s.include? 'query'
+      }
+      xml = Builder::XmlMarkup.new
+      xml.delete do
+        args.each do |spec|
+          spec.each do |k,v|
+            xml.tag!(k.to_sym, v)
+          end
+        end
+      end
+      http do |h|
+        h.update_search_index(index, xml.target!)
+      end
+      true
+    end
+    alias :delete_doc :remove
+    alias :deindex :remove
+  end
+end