lib/vertica/query.rb in vertica-0.9.1 vs lib/vertica/query.rb in vertica-0.9.2

@@ -1,37 +1,29 @@
 class Vertica::Query
 
-  attr_reader :connection, :sql
+  attr_reader :connection, :sql, :result, :error
   attr_accessor :row_handler, :copy_handler, :row_style
 
   def initialize(connection, sql, options = {})
     @connection, @sql = connection, sql
     
     @row_style    = options[:row_style] || @connection.row_style || :hash
     @row_handler  = options[:row_handler] 
     @copy_handler = options[:copy_handler]
+
+    @error  = nil
+    @result = Vertica::Result.new(row_style)
   end
   
   
   def run
-    @connection.write Vertica::Messages::Query.new(@sql)
-    result, error = nil, nil
+    @connection.write Vertica::Messages::Query.new(sql)
+    
     begin
-      case message = @connection.read_message
-      when Vertica::Messages::ErrorResponse
-        error = Vertica::Error::QueryError.from_error_response(message, @sql)
-      when Vertica::Messages::EmptyQueryResponse
-        error = Vertica::Error::EmptyQueryError.new("A SQL string was expected, but the given string was blank or only contained SQL comments.")
-      when Vertica::Messages::CopyInResponse
-        handle_copy_from_stdin
-      when Vertica::Messages::RowDescription, Vertica::Messages::CommandComplete
-        result = retreive_result(message, Vertica::Result.new(row_style))
-      else
-        @connection.process_message(message)
-      end
+      process_message(message = @connection.read_message)
     end until message.kind_of?(Vertica::Messages::ReadyForQuery)
-    
+
     raise error unless error.nil?
     return result
   end
   
   def write(data)
@@ -41,10 +33,29 @@
   
   alias_method :<<, :write
   
   protected
   
+  def process_message(message)
+    case message
+    when Vertica::Messages::ErrorResponse
+      @error = Vertica::Error::QueryError.from_error_response(message, @sql)
+    when Vertica::Messages::EmptyQueryResponse
+      @error = Vertica::Error::EmptyQueryError.new("A SQL string was expected, but the given string was blank or only contained SQL comments.")
+    when Vertica::Messages::CopyInResponse
+      handle_copy_from_stdin
+    when Vertica::Messages::RowDescription
+      result.descriptions = message
+    when Vertica::Messages::DataRow
+      handle_datarow(message)
+    when Vertica::Messages::CommandComplete
+      result.tag = message.tag
+    else
+      @connection.process_message(message)
+    end
+  end
+  
   def handle_copy_from_stdin
     if copy_handler.nil?
       @connection.write Vertica::Messages::CopyFail.new('no handler provided')
     else
       begin
@@ -57,28 +68,16 @@
         @connection.write Vertica::Messages::CopyFail.new(e.message)
         raise
       end
     end
   end
-  
-  def retreive_result(message, result)
-    until message.kind_of?(Vertica::Messages::CommandComplete)
-      case message
-      when Vertica::Messages::RowDescription
-        result.descriptions = message
-      when Vertica::Messages::DataRow
-        record = result.format_row(message)
-        result.add_row(record) if buffer_rows?
-        @row_handler.call(record) if @row_handler
-      else
-        @connection.process_message(message)
-      end
-      message = @connection.read_message
-    end
-    result.tag = message.tag
-    return result
+
+  def handle_datarow(datarow_message)
+    record = result.format_row(datarow_message)
+    result.add_row(record) if buffer_rows?
+    row_handler.call(record) if row_handler
   end
   
   def buffer_rows?
-    @row_handler.nil? && @copy_handler.nil?
+    row_handler.nil? && copy_handler.nil?
   end
 end