lib/pg_csv.rb in pg_csv-0.1.3 vs lib/pg_csv.rb in pg_csv-0.1.4

@@ -5,34 +5,42 @@
   def initialize(opts = {})
     @options = opts.symbolize_keys
   end
   
   # do export :to - filename or stream  
-  def export(to = nil, opts = {}, &block)
-    @block = block || Proc.new{|x|x}
+  def export(to = nil, opts = {}, &row_proc)
+    @row_proc = row_proc
     @local_options = opts.symbolize_keys
     
     raise ":connection should be" unless connection
     raise ":sql should be" unless sql
 
-    with_temp_file(to, temp_file, temp_dir) do |_to|
-      export_to(_to)
+    with_temp_file?(to, temp_file, temp_dir) do |dest|
+      export_to(dest)
     end        
   end
   
 protected
 
-  def with_temp_file(to, use_temp_file, tmp_dir)
+  def self.with_temp_file(dest, tmp_dir = '/tmp', &block)
+    require 'fileutils'
+    require 'tempfile'
+                                  
+    tempfile = Tempfile.new("pg_csv", tmp_dir)
+    yield(tempfile.path)
+        
+    FileUtils.mv(tempfile.path, dest)
+  end
+
+  def with_temp_file?(to, use_temp_file, tmp_dir)
     if use_temp_file && [:file, :gzip].include?(type)
       check_str(to)
+
+      self.class.with_temp_file(to, tmp_dir) do |filename|
+        yield(filename)
+      end
       
-      require 'fileutils'
-      require 'tempfile'
-                    
-      tempfile = Tempfile.new("pg_csv", tmp_dir)
-      yield(tempfile.path)
-      FileUtils.mv(tempfile.path, to)
       info "<=== moving export to #{to}"
     else
       yield(to)
     end
   end                                                            
@@ -64,11 +72,11 @@
         exporter[sio]
         result = sio.string
         
       when :yield
         # not real saving anywhere, just yield each record
-        raise "block should be" unless @block
+        raise "row_proc should be" unless @row_proc
         result = load_data{|_|}
     end
     
     info "<=== finished write #{to} in #{Time.now - start}"
     
@@ -103,15 +111,24 @@
     info "=> query"
     q = raw.exec(query)
     info "<= query"
 
     info "=> write data"
-    yield(@block[columns_str]) if columns_str
+    if columns_str
+      yield(@row_proc ? @row_proc[columns_str] : columns_str)
+    end
 
     count = 0    
-    while row = raw.get_copy_data()
-      yield(@block[row])
-      count += 1
+    if @row_proc
+      while row = raw.get_copy_data()
+        yield(@row_proc[row])
+        count += 1
+      end
+    else
+      while row = raw.get_copy_data()
+        yield(row)
+        count += 1
+      end
     end
     info "<= write data"
 
     q.clear
     count
\ No newline at end of file