cursor.rb in impala-0.2.0

- old
+ new

@@ -1,21 +1,23 @@
 module Impala
   # Cursors are used to iterate over result sets without loading them all
   # into memory at once. This can be useful if you're dealing with lots of
   # rows. It implements Enumerable, so you can use each/select/map/etc.
   class Cursor
+    BUFFER_SIZE = 1024
     include Enumerable
 
-    def initialize(handle, service, buffer_length=1024)
+    def initialize(handle, service)
       @handle = handle
       @service = service
 
-      @buffer_length = buffer_length
-      @row_buffer = []
 
+      @row_buffer = []
       @done = false
       @open = true
+
+      fetch_more
     end
 
     def inspect
       "#<#{self.class}#{open? ? '' : ' (CLOSED)'}>"
     end
@@ -73,13 +75,17 @@
     def metadata
       @metadata ||= @service.get_results_metadata(@handle)
     end
 
     def fetch_more
+      fetch_batch until @done || @row_buffer.count >= BUFFER_SIZE
+    end
+
+    def fetch_batch
       return if @done
 
       begin
-        res = @service.fetch(@handle, false, @buffer_length)
+        res = @service.fetch(@handle, false, BUFFER_SIZE)
       rescue Protocol::Beeswax::BeeswaxException => e
         @closed = true
         raise CursorError.new("Cursor has expired or been closed")
       end
 