Make downloading CSV work

diff --git a/src/pgwui_sql/views/sql.py b/src/pgwui_sql/views/sql.py
index 7e0d39851600e6f338e1cd3bb857989b3c787dd0..c169cea023807b4a27127ed624e43191d698f607 100644 (file)
--- a/src/pgwui_sql/views/sql.py
+++ b/src/pgwui_sql/views/sql.py
@@ -19,9 +19,14 @@
 
 from pyramid.view import view_config
 import attrs
+import codecs
+import csv
+import datetime
 import logging
 import markupsafe
 import psycopg.errors
+import pyramid.response
+import tempfile
 
 import pgwui_core.core
 import pgwui_core.utils
@@ -29,6 +34,8 @@ import pgwui_sql.views.base
 from pgwui_common.view import auth_base_view
 
 from pgwui_sql import exceptions as sql_ex
+from pgwui_core.constants import CSV
+from pgwui_sql.constants import MANY_FILES
 
 log = logging.getLogger(__name__)
 
@@ -79,6 +86,7 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
     sql_results = attrs.field(factory=list)
     search_path = attrs.field(default=None)     # requested search_path
     db_search_path = attrs.field(default=None)  # search_path of db
+    tfile = attrs.field(default=None)
 
     def make_form(self):
         return pgwui_sql.views.base.SQLForm().build(
@@ -89,6 +97,14 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
         super().read()
         self.search_path = self.request.POST.get('search_path')
 
+    def dl_filename(self):
+        uf = self.uf
+        return '_'.join(
+            ['sql_results',
+             uf["db"],
+             uf["user"],
+             datetime.datetime.now().isoformat(sep="_", timespec="seconds")])
+
     def write(self, result, errors):
         '''
         Setup dict to render resulting html form
@@ -105,6 +121,8 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
         response['filename'] = self.uf['filename']
         response['report_success'] = (not response['errors']
                                       and self.uf['action'] != '')
+        if self.uf['download']:
+            response['dl_filename'] = self.dl_filename()
         return response
 
     def val_input(self):
@@ -181,11 +199,25 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
             return sql
         return f'SET search_path TO {self.search_path};\n{sql}'
 
+    def make_download(self, sql_results):
+        downloading = self.uf['download']
+        if downloading:
+            self.tfile = tempfile.TemporaryFile(mode='w+t', newline='')
+            writer = csv.writer(self.tfile)
+            for sql_result in sql_results:
+                writer.writerow(sql_result.statusmessage.data)
+                writer.writerow(sql_result.heading.data)
+                for row in sql_result.rows:
+                    writer.writerow(row.data)
+                writer.writerow(f'{sql_result.rowcount.data} rows')
+            self.tfile.seek(0)
+
     def cleanup(self):
         '''
         Execute a series of SQL statements.
-        The result goes into the upload handler (uh.sql_results),
-        interleaving errors with output.
+        The result goes into the upload handler (uh.sql_results).
+        This happens even when a file is built to download so that
+        should there be errors the results that do exist are displayed.
         '''
         cur = self.cur
         null_rep = self.uf['null_rep']
@@ -213,6 +245,8 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler):
             sql_results.append(sql_result)
             nextset = cur.nextset()
 
+        self.make_download(sql_results)
+
     def factory(self, ue):
         '''Make a db loader function from an UploadEngine.
 
@@ -234,14 +268,14 @@ def log_response(response):
         else:
             download_fmt = 'TAB'
         if pgwui_core.utils.is_checked(response['one_file_checked']):
-            download_as = 'One file'
+            download_as = 'one file'
         else:
-            download_as = 'A zip file of one file per result set'
+            download_as = 'a zip file of one file per result set'
         log.info(': '.join(
             [f'Successful sql: DB {response["db"]}',
              # ' File ({filename}):'
              f'Download Format {download_fmt}',
-             f'Download AS {download_as}',
+             f'Download as {download_as}',
              f'By user {response["user"]}']))
 
 
@@ -253,11 +287,28 @@ def sql_view(request):
     uh = SQLResultsHandler(request).init()
     response = pgwui_core.core.UploadEngine(uh).run()
 
-    settings = request.registry.settings
-    response.setdefault('pgwui', dict())
-    response['pgwui']['pgwui_sql'] = settings['pgwui']['pgwui_sql']
+    if uh.uf['download'] and not response['errors']:
+        pmd_response = pyramid.response.Response()
+        pmd_response.cache_control = 'public; max-age=0'
+        if uh.uf['download_as'] == MANY_FILES:
+            pmd_response.content_type = 'application/zip'
+        elif uh.uf['download_fmt'] == CSV:
+            pmd_response.content_type = 'text/csv'
+        else:
+            pmd_response.content_type = 'text/plain'
+        pmd_response.content_disposition = \
+            f'attachment; filename={response["dl_filename"]}.csv'
+
+        pmd_response.app_iter = codecs.iterencode(uh.tfile, 'utf_8')
 
-    response['result_rows'] = uh.sql_results
+        log_response(response)
+        return pmd_response
+    else:
+        settings = request.registry.settings
+        response.setdefault('pgwui', dict())
+        response['pgwui']['pgwui_sql'] = settings['pgwui']['pgwui_sql']
 
-    log_response(response)
-    return response
+        response['result_rows'] = uh.sql_results
+
+        log_response(response)
+        return response

diff --git a/tests/views/test_sql.py b/tests/views/test_sql.py
index 00625deb7aefb0fa3ffd33027e170b04590c7797..f54b85bf1800bb454f2e1ec9f425342e06583705 100644 (file)
--- a/tests/views/test_sql.py
+++ b/tests/views/test_sql.py
@@ -53,9 +53,13 @@ CHANGED_RESPONSE = {
     'trim_upload': constants.CHECKED,
     'upload_null': constants.CHECKED,
     'user': 'someuser',
+    'download': False,
 }
 
-UNCHANGED_RESPONSE = {'report_success': False}
+UNCHANGED_RESPONSE = {
+    'report_success': False,
+    'download': False,
+}
 
 
 # Helper classes
@@ -78,11 +82,14 @@ def isolate_sql_view(monkeypatch, pyramid_request_config):
     cause UploadEngine to return the supplied "response".
     '''
     def run(response):
-        def upload_engine(*args):
+        def upload_engine(uh, *args):
+            uh.read()
             return MockUploadEngine(response)
 
         monkeypatch.setattr(
             sql.pgwui_core.core, 'UploadEngine', upload_engine)
+        monkeypatch.setattr(
+            sql, 'log_response', lambda *args: None)
 
         settings = pyramid_request_config.get_settings()
         settings['pgwui'] = settings.get('pgwui', dict())