Add support for scrolling & efficient index copying.

Author: jsuchal

lib/es.rb

@@ -1,10 +1,11 @@
-require 'es/version'
-require 'es/connection'
-require 'es/raw_client'
-require 'es/client'
-
 module ES
   def self.new(*args) # faux constructor
     Client.new(*args)
   end
 end
+
+require 'es/version'
+require 'es/connection'
+require 'es/raw_client'
+require 'es/client'
+require 'es/utils'

lib/es/client.rb

@@ -29,9 +29,14 @@ def index(path, data)
       @dumper.load(response)
     end
 
-    def search(path, data)
+    def search(path, data, params = {})
       serialized = serialize(data)
-      response = @client.search(path, serialized)
+      response = @client.search(path, serialized, params)
+      @dumper.load(response)
+    end
+
+    def scroll(params = {})
+      response = @client.scroll(params)
       @dumper.load(response)
     end
 
@@ -52,6 +57,11 @@ def delete_index(path)
       @dumper.load(response)
     end
 
+    def get_mapping(path)
+      response = @client.get_mapping(path)
+      @dumper.load(response)
+    end
+
     private
     def serialize(data)
       @dumper.dump(data, mode: :compat)

lib/es/raw_client.rb

@@ -24,8 +24,12 @@ def index(path, data)
       @connection.request(:put, path, data)
     end
 
-    def search(path, query)
-      @connection.request(:post, action_path(path, :search), query)
+    def search(path, query, params = {})
+      @connection.request(:post, action_path(path, :search, params), query)
+    end
+
+    def scroll(params = {})
+      @connection.request(:get, action_path(nil, 'search/scroll', params))
     end
 
     def update(path, data)
@@ -36,9 +40,18 @@ def bulk(requests, path = nil)
       @connection.request(:post, action_path(path, :bulk), requests)
     end
 
+    def get_mapping(path)
+      @connection.request(:get, action_path(path, :mapping))
+    end
+
     private
-    def action_path(path, action)
-      path ? "#{path}/_#{action}" : "_#{action}"
+    def action_path(path, action, params = {})
+      full_path = path ? "#{path}/_#{action}" : "_#{action}"
+      params.any? ? "#{full_path}?#{serialize_params(params)}" : full_path
+    end
+
+    def serialize_params(params)
+      params.map { |k, v| "#{k}=#{v}" }.join('&')
     end
   end
 end

lib/es/utils.rb

@@ -0,0 +1,23 @@
+module ES
+  def self.copy_index(from_index, to_index, opts = {})
+    scroll_timeout = opts.fetch(:scroll_timeout, '5m')
+    scroll_size = opts.fetch(:scroll_size, 100)
+
+    from = new(host: opts[:from_host])
+    to = new(host: opts[:to_host])
+
+    results = from.search(from_index, {query: {match_all: {}}}, {scroll: scroll_timeout, size: scroll_size})
+
+    while results['hits']['hits'].any?
+      requests = results['hits']['hits'].each_with_object([]) do |hit, requests|
+        requests << {index: {_type: hit['_type'], _id: hit['id']}}
+        requests << hit['_source']
+      end
+
+      to.bulk(requests, to_index)
+
+      scroll_id = results['_scroll_id']
+      results = from.scroll(scroll: scroll_timeout, scroll_id: scroll_id)
+    end
+  end
+end

spec/es/client_spec.rb

@@ -42,14 +42,32 @@
   end
 
   it 'serializes data for .search' do
-    client.should_receive(:search).with('index/1', raw_data).and_return(raw_response)
+    client.should_receive(:search).with('index/1', raw_data, {size: 20}).and_return(raw_response)
 
-    subject.search('index/1', data).should == response
+    subject.search('index/1', data, size: 20).should == response
+  end
+
+  it 'serializes data from .scroll' do
+    client.should_receive(:scroll).with({scroll_id: 123}).and_return(raw_response)
+
+    subject.scroll(scroll_id: 123)
   end
 
   it 'serializes data for .update' do
     client.should_receive(:update).with('index/1', raw_data).and_return(raw_response)
 
     subject.update('index/1', data).should == response
   end
+
+  it 'unserializes response for .get_mapping' do
+    client.should_receive(:get_mapping).with('index').and_return(raw_response)
+
+    subject.get_mapping('index').should == response
+  end
+
+  it 'should respond to all methods of raw client' do
+    ES::RawClient.instance_methods(false).each do |method|
+      subject.respond_to?(method).should be_true, "expected to respond to message #{method}"
+    end
+  end
 end

spec/es/connection_spec.rb

@@ -3,7 +3,7 @@
 describe ES::Connection do
   let(:data) { :data }
   let(:response_body) { :response }
-  let(:response) { stub(response_code: 200, body_str: response_body) }
+  let(:response) { double(response_code: 200, body_str: response_body) }
   let(:driver) { double(:Driver) }
 
   subject { described_class.new('http://localhost:9200', driver) }
@@ -33,7 +33,7 @@
   end
 
   it 'throws error if response code is not 2XX' do
-    error_response = stub(response_code: 400)
+    error_response = double(response_code: 400)
     driver.should_receive(:http_post).and_return(error_response)
 
     expect { subject.request(:post, 'index/_search', data) }.to raise_error(ES::Connection::Error)

spec/es/raw_client_spec.rb

@@ -44,6 +44,20 @@
 
       subject.search('index', data).should == response
     end
+
+    it 'sends POST to connection with parameters' do
+      connection.should_receive(:request).with(:post, 'index/_search?scroll=5m&size=100', data).and_return(response)
+
+      subject.search('index', data, scroll: '5m', size: 100)
+    end
+  end
+
+  context '.scroll' do
+    it 'sends GET to connection with parameters' do
+      connection.should_receive(:request).with(:get, '_search/scroll?scroll=5m&scroll_id=123').and_return(response)
+
+      subject.scroll(scroll: '5m', scroll_id: 123)
+    end
   end
 
   context '.bulk' do
@@ -61,4 +75,12 @@
       subject.update('index/1', data).should == response
     end
   end
+
+  context '.get_mapping' do
+    it 'sends GET to connection' do
+      connection.should_receive(:request).with(:get, 'index/_mapping').and_return(response)
+
+      subject.get_mapping('index').should == response
+    end
+end
 end