erickt · March 29, 2012 17:49 · Mar 29, 2012
diff --git a/gistfile1.rb b/gistfile1.rb
@@ -0,0 +1,122 @@
+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