Create a gist now

Instantly share code, notes, and snippets.

module Tire
module Scroll
DEFAULT_SCROLL_DURATION = '1m'
class Scroll
attr_reader :search, :total, :remaining, :counter
def initialize(indices=nil, options={}, &block)
@search = Tire::Search::Search.new(indices, options, &block)
@options = options
@duration = options[:scroll] || DEFAULT_SCROLL_DURATION
end
def url
Configuration.url + "/_search/scroll?scroll=#{@duration}"
end
def results
@results || (perform; @results)
end
def response
@response || (perform; @response)
end
def each
begin
perform
results.each do |item|
pp ['each', item.id]
yield item
end
end while any?
end
def in_batches
begin
perform
yield @results
end while any?
end
def done?
!any?
end
def any?
@remaining > 0
end
def perform
if @response.nil?
puts @search.to_curl
puts
@response = @search.response
@json = @search.json
@results = @search.results
reset_stats
else
puts to_curl
puts
begin
@response = Configuration.client.post(url, @scroll_id)
if @response.failure?
STDERR.puts "[REQUEST FAILED] #{self.to_curl}\n"
raise SearchRequestFailed, @response.to_s
end
@json = MultiJson.decode(@response.body)
@results = Results::Collection.new(@json, @options)
update_stats
ensure
logged
end
end
@scroll_id = @json['_scroll_id']
return self
end
def to_curl
%Q|curl -X GET "#{url}?pretty=true" -d '#{@scroll_id}'|
end
def logged(error=nil)
if Configuration.logger
Configuration.logger.log_request 'scroll', nil, to_curl
took = @json['took'] rescue nil
code = @response.code rescue nil
if Configuration.logger.level.to_s == 'debug'
# FIXME: Depends on RestClient implementation
body = if @json
defined?(Yajl) ? Yajl::Encoder.encode(@json, :pretty => true) : MultiJson.encode(@json)
else
@response.body rescue nil
end
else
body = ''
end
Configuration.logger.log_response code || 'N/A', took || 'N/A', body || 'N/A'
end
end
protected
def update_stats
@remaining -= @results.count
@counter += @results.count
end
def reset_stats
@total = @json['hits']['total']
@counter = @json['hits']['hits'].length
@remaining = @total - @counter
end
end
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment