lowang-rubberband 0.0.3

data/lib/elasticsearch/encoding/json.rb

@@ -0,0 +1,19 @@
+require 'yajl'
+
+module ElasticSearch
+  module Encoding
+    class JSON < Base
+      def encode(object)
+        Yajl::Encoder.encode(object)
+      end
+
+      def decode(string)
+        Yajl::Parser.parse(string)
+      end
+
+      def is_encoded?(object)
+        object.is_a?(String)
+      end
+    end
+  end
+end

data/lib/elasticsearch/transport.rb

@@ -0,0 +1,14 @@
+require "transport/base_protocol"
+require "transport/base"
+
+module ElasticSearch
+  class ConnectionFailed < RetryableError; end
+  class HostResolutionError < RetryableError; end
+  class TimeoutError < RetryableError; end
+  class RequestError < FatalError; end
+
+  module Transport
+    autoload :HTTP, 'transport/http'
+    autoload :Thrift, 'transport/thrift'
+  end
+end

data/lib/elasticsearch/transport/base.rb

@@ -0,0 +1,37 @@
+module ElasticSearch
+  class RetryableError < StandardError; end
+  class FatalError < StandardError; end
+
+  module Transport
+
+    DEFAULTS = {
+      :encoder => ElasticSearch::Encoding::JSON
+    }.freeze
+
+    class Base
+      include BaseProtocol
+
+      attr_accessor :server, :options
+
+      def initialize(server, options={})
+        @server = server
+        @options = DEFAULTS.merge(options)
+      end
+
+      def connect!
+        raise NotImplementedError
+      end
+
+      def close
+      end
+
+      def encoder
+        @encoder ||= @options[:encoder].new
+      end
+
+      def request(method, operation, params={}, body=nil, headers={})
+        raise NotImplementedError
+      end
+    end
+  end
+end

data/lib/elasticsearch/transport/base_protocol.rb

@@ -0,0 +1,235 @@
+module ElasticSearch
+  module Transport
+    module IndexProtocol
+      def index(index, type, id, document, options={})
+        body = encoder.is_encoded?(document) ? document : encoder.encode(document)
+        if id.nil?
+          response = request(:post, {:index => index, :type => type}, {}, body)
+        else
+          response = request(:put, {:index => index, :type => type, :id => id}, {}, body)
+        end
+        handle_error(response) unless response.status == 200
+        encoder.decode(response.body)
+      end
+
+      def get(index, type, id, options={})
+        response = request(:get, {:index => index, :type => type, :id => id})
+        return nil if response.status == 404
+
+        handle_error(response) unless response.status == 200
+        hit = encoder.decode(response.body)
+        unescape_id!(hit) #TODO extract these two calls from here and search
+        set_encoding!(hit)
+        hit # { "_id", "_index", "_type", "_source" }
+      end
+
+      def delete(index, type, id, options={})
+        response = request(:delete,{:index => index, :type => type, :id => id})
+        handle_error(response) unless response.status == 200 # ElasticSearch always returns 200 on delete, even if the object doesn't exist
+        encoder.decode(response.body)
+      end
+
+      def search(index, type, query, options={})
+        if query.is_a?(Hash)
+          # patron cannot submit get requests with content, so if query is a hash, post it instead (assume a query hash is using the query dsl)
+          response = request(:post, {:index => index, :type => type, :op => "_search"}, options, encoder.encode(query))
+        else
+          response = request(:get, {:index => index, :type => type, :op => "_search"}, options.merge(:q => query))
+        end
+        handle_error(response) unless response.status == 200
+        results = encoder.decode(response.body)
+        # unescape ids
+        results["hits"]["hits"].each do |hit|
+          unescape_id!(hit)
+          set_encoding!(hit)
+        end
+        results # {"hits"=>{"hits"=>[{"_id", "_type", "_source", "_index", "_score"}], "total"}, "_shards"=>{"failed", "total", "successful"}}
+      end
+
+      def scroll(scroll_id)
+        response = request(:get, {:op => "_search/scroll"}, {:scroll_id => scroll_id })
+        handle_error(response) unless response.status == 200
+        results = encoder.decode(response.body)
+        # unescape ids
+        results["hits"]["hits"].each do |hit|
+          unescape_id!(hit)
+          set_encoding!(hit)
+        end
+        results # {"hits"=>{"hits"=>[{"_id", "_type", "_source", "_index", "_score"}], "total"}, "_shards"=>{"failed", "total", "successful"}, "_scrollId"}
+      end
+
+      def count(index, type, query, options={})
+        if query.is_a?(Hash)
+          # patron cannot submit get requests with content, so if query is a hash, post it instead (assume a query hash is using the query dsl)
+          response = request(:post, {:index => index, :type => type, :op => "_count"}, options, encoder.encode(query))
+        else
+          response = request(:get, {:index => index, :type => type, :op => "_count"}, options.merge(:q => query))
+        end
+        handle_error(response) unless response.status == 200
+        encoder.decode(response.body) # {"count", "_shards"=>{"failed", "total", "successful"}}
+      end
+    end
+
+    module IndexAdminProtocol
+      def index_status(index_list, options={})
+        standard_request(:get, {:index => index_list, :op => "_status"})
+      end
+
+      def create_index(index, create_options={}, options={})
+        standard_request(:put, {:index => index}, {}, encoder.encode(create_options))
+      end
+
+      def delete_index(index, options={})
+        standard_request(:delete, {:index => index})
+      end
+
+      def alias_index(operations, options={})
+        standard_request(:post, {:op => "_aliases"}, {}, encoder.encode(operations))
+      end
+
+      def update_mapping(index, type, mapping, options)
+        standard_request(:put, {:index => index, :type => type, :op => "_mapping"}, options, encoder.encode(mapping))
+      end
+
+      def flush(index_list, options={})
+        standard_request(:post, {:index => index_list, :op => "_flush"}, options, "")
+      end
+      
+      def refresh(index_list, options={})
+        standard_request(:post, {:index => index_list, :op => "_refresh"}, {}, "")
+      end
+      
+      def snapshot(index_list, options={})
+        standard_request(:post, {:index => index_list, :type => "_gateway", :op => "snapshot"}, {}, "")
+      end
+      
+      def optimize(index_list, options={})
+        standard_request(:post, {:index => index_list, :op => "_optimize"}, options, {}, "")
+      end
+    end
+
+    module ClusterAdminProtocol
+      def cluster_health(index_list, options={})
+        standard_request(:get, {:index => "_cluster", :type => "health", :id => index_list}, options)
+      end
+
+      def cluster_state(options={})
+        standard_request(:get, {:index => "_cluster", :op => "state"})
+      end
+
+      def nodes_info(node_list, options={})
+        standard_request(:get, {:index => "_cluster", :type => "nodes", :id => node_list})
+      end
+
+      def nodes_stats(node_list, options={})
+        standard_request(:get, {:index => "_cluster", :type => "nodes", :id => node_list, :op => "stats"})
+      end
+
+      def shutdown_nodes(node_list, options={})
+        standard_request(:post, {:index => "_cluster", :type => "nodes", :id => node_list, :op => "_shutdown"}, options, "")
+      end
+
+      def restart_nodes(node_list, options={})
+        standard_request(:post, {:index => "_cluster", :type => "nodes", :id => node_list, :op =>  "_restart"}, options, "")
+      end
+    end
+
+    module ProtocolHelpers
+      private
+
+      def standard_request(*args)
+        response = request(*args)
+        handle_error(response) unless response.status == 200
+        encoder.decode(response.body)
+      end
+
+      def handle_error(response)
+        raise RequestError, "(#{response.status}) #{response.body}"
+      end
+
+      # :index - one or many index names
+      # :type - one or many types
+      # :id - one id
+      # :op - optional operation
+      def generate_uri(options)
+        path = ""
+        path << "/#{Array(options[:index]).collect { |i| escape(i.downcase) }.join(",")}" if options[:index] && !options[:index].empty?
+        path << "/#{Array(options[:type]).collect { |t| escape(t) }.join(",")}" if options[:type] && !options[:type].empty?
+        path << "/#{Array(options[:id]).collect { |id| escape(id) }.join(",")}" if options[:id] && !options[:id].empty?
+        path << "/#{options[:op]}" if options[:op]
+        path
+      end
+
+      #doesn't handle arrays or hashes or what have you
+      def generate_query_string(params)
+        params.collect { |k,v| "#{escape(k.to_s)}=#{escape(v.to_s)}" }.join("&")
+      end
+
+      def unescape_id!(hit)
+        hit["_id"] = unescape(hit["_id"])
+        nil
+      end
+
+      def set_encoding!(hit)
+        encode_utf8(hit["_source"]) if hit["_source"].is_a?(String)
+        nil
+      end
+
+      # faster than CGI.escape
+      # stolen from RSolr, who stole it from Rack
+      def escape(string)
+        string.to_s.gsub(/([^ a-zA-Z0-9_.-]+)/n) {
+          #'%'+$1.unpack('H2'*$1.size).join('%').upcase
+          '%'+$1.unpack('H2'*bytesize($1)).join('%').upcase
+        }.tr(' ', '+')
+      end
+
+      def unescape(string)
+        CGI.unescape(string)
+      end
+
+      if ''.respond_to?(:force_encoding) && ''.respond_to?(:encoding)
+        # encodes the string as utf-8 in Ruby 1.9
+        def encode_utf8(string)
+          # ElasticSearch only ever returns json in UTF-8 (per the JSON spec) so we can use force_encoding here (#TODO what about ids? can we assume those are always ascii?)
+          string.force_encoding(::Encoding::UTF_8)
+        end
+      else
+        # returns the unaltered string in Ruby 1.8
+        def encode_utf8(string)
+          string
+        end
+      end
+
+      # Return the bytesize of String; uses String#size under Ruby 1.8 and
+      # String#bytesize under 1.9.
+      if ''.respond_to?(:bytesize)
+        def bytesize(string)
+          string.bytesize
+        end
+      else
+        def bytesize(string)
+          string.size
+        end
+      end
+    end
+
+    module BaseProtocol
+      include IndexProtocol
+      include IndexAdminProtocol
+      include ClusterAdminProtocol
+      include ProtocolHelpers
+
+      def all_nodes
+        http_addresses = nodes_info([])["nodes"].collect { |id, node| node["http_address"] }
+        http_addresses.collect! do |a|
+          if a =~ /inet\[.*\/([\d.:]+)\]/
+            $1
+          end
+        end.compact!
+        http_addresses
+      end
+
+    end
+  end
+end

data/lib/elasticsearch/transport/http.rb

@@ -0,0 +1,50 @@
+require 'patron'
+require 'cgi'
+
+module ElasticSearch
+  module Transport
+    class HTTP < Base
+
+      DEFAULTS = {
+        :timeout => 5
+      }.freeze
+
+      def initialize(server, options={})
+        super
+        @options = DEFAULTS.merge(@options)
+      end
+
+      def connect!
+        @session = Patron::Session.new
+        @session.base_url = @server
+        @session.timeout = @options[:timeout]
+        @session.headers['User-Agent'] = 'ElasticSearch.rb v0.1'
+      end
+
+      private
+
+      def request(method, operation, params={}, body=nil, headers={})
+        begin
+          uri = generate_uri(operation)
+          query = generate_query_string(params)
+          path = [uri, query].join("?")
+          #puts "request: #{@server} #{path} #{body}"
+          response = @session.request(method, path, headers, :data => body)
+          handle_error(response) if response.status >= 500
+          response
+        rescue Exception => e
+          case e
+          when Patron::ConnectionFailed
+            raise ConnectionFailed
+          when Patron::HostResolutionError
+            raise HostResolutionError
+          when Patron::TimeoutError
+            raise TimeoutError
+          else
+            raise e
+          end
+        end
+      end
+    end
+  end
+end

data/lib/elasticsearch/transport/thrift.rb

@@ -0,0 +1,96 @@
+require 'thrift'
+$LOAD_PATH.unshift(File.join(File.dirname(__FILE__), "thrift"))
+require 'rest'
+
+module ElasticSearch
+  module Transport
+    class Thrift < Base
+
+      DEFAULTS = {
+        :timeout => 5,
+        :thrift_protocol => ::Thrift::BinaryProtocol,
+        :thrift_protocol_extra_params => [],
+        :thrift_transport => ::Thrift::Socket,
+        :thrift_transport_wrapper => ::Thrift::BufferedTransport,
+        :client_class => ElasticSearch::Thrift::Rest::Client
+      }.freeze
+
+      def initialize(server, options={})
+        super
+        @options = DEFAULTS.merge(@options)
+      end
+
+      def connect!
+        host, port = parse_server(@server)
+
+        @transport = @options[:thrift_transport].new(host, port.to_i, @options[:timeout])
+        @transport = @transport_wrapper.new(@transport) if @transport_wrapper
+        @transport.open
+
+        @client = @options[:client_class].new(@options[:thrift_protocol].new(@transport, *@options[:thrift_protocol_extra_params]))
+      rescue ::Thrift::TransportException, Errno::ECONNREFUSED
+        close
+        raise ConnectionFailed
+      end
+
+      def close
+        @transport.close rescue nil
+      end
+
+      private
+
+      def parse_server(server)
+        host, port = server.to_s.split(":")
+        raise ArgumentError, 'Servers must be in the form "host:port"' unless host and port
+        [host, port]
+      end
+
+      def stringify!(hash)
+        hash.keys.each do |k|
+          hash[k.to_s] = hash.delete(k).to_s
+        end
+        hash
+      end
+
+      def request(method, operation, params={}, body=nil, headers={})
+        begin
+          uri = generate_uri(operation)
+          #puts "request: #{@server} #{method} #{uri} #{params.inspect} #{body}"
+          request = ElasticSearch::Thrift::RestRequest.new
+          case method
+          when :get
+            request.method = ElasticSearch::Thrift::Method::GET
+          when :put
+            request.method = ElasticSearch::Thrift::Method::GET
+          when :post
+            request.method = ElasticSearch::Thrift::Method::POST
+          when :delete
+            request.method = ElasticSearch::Thrift::Method::DELETE
+          end
+
+          request.uri = uri
+          request.params = stringify!(params) #TODO this will change to parameters= in versions > 0.11.0
+          request.body = body
+          request.headers = stringify!(headers)
+          response = @client.execute(request)
+          handle_error(response) if response.status >= 500
+          response
+        rescue Exception => e
+          case e
+          when ::Thrift::TransportException
+            case e.type
+            when ::Thrift::TransportException::TIMED_OUT
+              raise TimeoutError
+            else
+              raise ConnectionFailed
+            end
+          #TODO Thrift::ApplicationException, Thrift::ProtocolException, IOError.. retryable or fatal?
+          else
+            raise e
+          end
+        end
+      end
+
+    end
+  end
+end