From 01e74f53a1f7ebd8309c41a5d89969a48fb3e5f2 Mon Sep 17 00:00:00 2001
From: Graham Jackson <graham.jackson07@gmail.com>
Date: Wed, 15 Apr 2015 12:30:46 -0400
Subject: [PATCH] Add option to include headers for #to_csv, and #to_tsv

---
 lib/rbhive/result_set.rb | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/lib/rbhive/result_set.rb b/lib/rbhive/result_set.rb
index c5f3d61..2b2f820 100644
--- a/lib/rbhive/result_set.rb
+++ b/lib/rbhive/result_set.rb
@@ -4,34 +4,39 @@ def initialize(rows, schema)
       @schema = schema
       super(rows.map {|r| @schema.coerce_row(r) })
     end
-    
+
     def column_names
       @schema.column_names
     end
-    
+
     def column_type_map
       @schema.column_type_map
     end
-    
-    def to_csv(out_file=nil)
-      to_seperated_output(",", out_file)
+
+    def to_csv(out_file=nil, opts={})
+      to_seperated_output(",", out_file, opts)
     end
-    
-    def to_tsv(out_file=nil)
-      to_seperated_output("\t", out_file)
+
+    def to_tsv(out_file=nil, opts={})
+      to_seperated_output("\t", out_file, opts)
     end
-    
+
     def as_arrays
       @as_arrays ||= self.map{ |r| @schema.coerce_row_to_array(r) }
     end
-    
+
     private
-    
-    def to_seperated_output(sep, out_file)
+
+    def to_seperated_output(sep, out_file, opts)
       rows = self.map { |r| @schema.coerce_row_to_array(r).join(sep) }
+      rows.insert(0, separated_headers(sep)) if opts[:headers]
       sv = rows.join("\n")
       return sv if out_file.nil?
       File.open(out_file, 'w+') { |f| f << sv }
     end
+
+    def separated_headers(sep)
+      column_names.join(sep)
+    end
   end
-end
\ No newline at end of file
+end