From cfa3cc7c59bc0bde7c0edfc4772e4edf749ed078 Mon Sep 17 00:00:00 2001
From: Toms Medins <toms.medins@makit.lv>
Date: Wed, 5 Jun 2019 13:58:29 +0300
Subject: [PATCH] Added `to_pgxlsx` and `to_pgcsv` support

---
 lib/xport/formatters/axlsx.rb | 30 ++++++++++++++++++++++++++++++
 lib/xport/formatters/csv.rb   | 25 +++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/lib/xport/formatters/axlsx.rb b/lib/xport/formatters/axlsx.rb
index 7d34ff5..7f1ea69 100644
--- a/lib/xport/formatters/axlsx.rb
+++ b/lib/xport/formatters/axlsx.rb
@@ -13,6 +13,36 @@ def to_xlsx(&block)
       to_file(formatter, &block)
     end
 
+    def to_pgxlsx
+      formatter = Xport::Axlsx::Formatter.new(self)
+      formatter.add_worksheet do |worksheet|
+        write_header(formatter, worksheet)
+        connection = object_class.connection
+        query = connection.unprepared_statement do
+          scope = object_class.from("(#{objects.to_sql}) AS #{object_class.table_name}")
+          builder.columns.each_with_index do |column, i|
+            name = builder.headers[i]
+            as = human_attribute_name(name)
+            scope = scope.select("#{column} AS \"#{as}\"")
+          end
+          "COPY (#{scope.to_sql}) TO STDOUT"
+        end
+        raw_connection = connection.raw_connection
+        connection.send(:log, query) do
+          deco = PG::TextDecoder::CopyRow.new
+          res = raw_connection.copy_data query, deco do
+            while row = raw_connection.get_copy_data
+              formatter.add_row worksheet, row
+            end
+          end
+          res.check
+        end
+
+        write_widths(formatter, worksheet)
+      end
+      formatter.to_file
+    end
+
     class Formatter
       attr_reader :export, :workbook
 
diff --git a/lib/xport/formatters/csv.rb b/lib/xport/formatters/csv.rb
index a1401cc..59cd7d2 100644
--- a/lib/xport/formatters/csv.rb
+++ b/lib/xport/formatters/csv.rb
@@ -13,6 +13,31 @@ def to_csv(&block)
       to_file(formatter, &block)
     end
 
+    def to_pgcsv
+      StringIO.new.tap do |io|
+        connection = object_class.connection
+        query = connection.unprepared_statement do
+          scope = object_class.from("(#{objects.to_sql}) AS #{object_class.table_name}")
+          builder.columns.each_with_index do |column, i|
+            name = builder.headers[i]
+            as = human_attribute_name(name)
+            scope = scope.select("#{column} AS \"#{as}\"")
+          end
+          "COPY (#{scope.to_sql}) TO STDOUT WITH CSV HEADER"
+        end
+        raw_connection = connection.raw_connection
+        connection.send(:log, query) do
+          res = raw_connection.copy_data query do
+            while row = raw_connection.get_copy_data
+              io << row
+            end
+          end
+          res.check
+        end
+        io.rewind
+      end
+    end
+
     class Formatter
       def initialize(export)
         @io  = StringIO.new