gistfile1.rb

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