Fix various issues in CSV file download feature:

1) To handle non-ascii filenames which we set from table name. Fixes #2314

2) To handle non-ascii query data. Fixes #2253

3) To dump JSON type columns properly in csv. Fixes #2360

requirements.txt

@@ -40,3 +40,4 @@ SQLAlchemy==1.0.14
 sqlparse==0.1.19
 Werkzeug==0.9.6
 WTForms==2.0.2
+backports.csv==1.0.4; python_version <= '2.7'

web/pgadmin/tools/sqleditor/__init__.py

@@ -1337,7 +1337,8 @@ def save_file():
 @login_required
 def start_query_download_tool(trans_id):
     sync_conn = None
-    status, error_msg, conn, trans_obj, session_obj = check_transaction_status(trans_id)
+    status, error_msg, conn, trans_obj, \
+        session_obj = check_transaction_status(trans_id)
 
     if status and conn is not None \
             and trans_obj is not None and session_obj is not None:
@@ -1361,11 +1362,15 @@ def start_query_download_tool(trans_id):
                     del conn.manager.connections[sync_conn.conn_id]
 
                 # This returns generator of records.
-                status, gen = sync_conn.execute_on_server_as_csv(sql, records=2000)
+                status, gen = sync_conn.execute_on_server_as_csv(
+                    sql, records=2000
+                )
 
                 if not status:
                     r = Response('"{0}"'.format(gen), mimetype='text/csv')
-                    r.headers["Content-Disposition"] = "attachment;filename=error.csv"
+                    r.headers[
+                        "Content-Disposition"
+                    ] = "attachment;filename=error.csv"
                     r.call_on_close(cleanup)
                     return r
 
@@ -1377,7 +1382,18 @@ def start_query_download_tool(trans_id):
                     import time
                     filename = str(int(time.time())) + ".csv"
 
-                r.headers["Content-Disposition"] = "attachment;filename={0}".format(filename)
+                # We will try to encode report file name with latin-1
+                # If it fails then we will fallback to default ascii file name
+                # werkzeug only supports latin-1 encoding supported values
+                try:
+                    tmp_file_name = filename
+                    tmp_file_name.encode('latin-1', 'strict')
+                except UnicodeEncodeError:
+                    filename = "download.csv"
+
+                r.headers[
+                    "Content-Disposition"
+                ] = "attachment;filename={0}".format(filename)
 
                 r.call_on_close(cleanup)
                 return r
@@ -1388,4 +1404,6 @@ def start_query_download_tool(trans_id):
             r.call_on_close(cleanup)
             return r
     else:
-        return internal_server_error(errormsg=gettext("Transaction status check failed."))
+        return internal_server_error(
+            errormsg=gettext("Transaction status check failed.")
+        )

web/pgadmin/utils/driver/psycopg2/__init__.py

@@ -18,8 +18,8 @@ import os
 import random
 import select
 import sys
-import csv
 
+import simplejson as json
 import psycopg2
 import psycopg2.extras
 from flask import g, current_app, session
@@ -36,11 +36,15 @@ from ..abstract import BaseDriver, BaseConnection
 from .cursor import DictCursor
 
 if sys.version_info < (3,):
+    # Python2 in-built csv module do not handle unicode
+    # backports.csv module ported from PY3 csv module for unicode handling
+    from backports import csv
     from StringIO import StringIO
     psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
     psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
 else:
     from io import StringIO
+    import csv
 
 _ = gettext
 
@@ -597,7 +601,22 @@ WHERE
         if self.async == 1:
             self._wait(cur.connection)
 
-    def execute_on_server_as_csv(self, query, params=None, formatted_exception_msg=False, records=2000):
+    def execute_on_server_as_csv(self,
+                                 query, params=None,
+                                 formatted_exception_msg=False,
+                                 records=2000):
+        """
+        To fetch query result and generate CSV output
+
+        Args:
+            query: SQL
+            params: Additional parameters
+            formatted_exception_msg: For exception
+            records: Number of initial records
+
+        Returns:
+            Generator response
+        """
         status, cur = self.__cursor(server_cursor=True)
         self.row_count = 0
 
@@ -608,21 +627,26 @@ WHERE
         if sys.version_info < (3,) and type(query) == unicode:
             query = query.encode('utf-8')
 
-        current_app.logger.log(25,
+        current_app.logger.log(
-                               u"Execute (with server cursor) for server #{server_id} - {conn_id} (Query-id: {query_id}):\n{query}".format(
+           25,
-                                   server_id=self.manager.sid,
+           u"Execute (with server cursor) for server #{server_id} - {conn_id} "
-                                   conn_id=self.conn_id,
+           u"(Query-id: {query_id}):\n{query}".format(
-                                   query=query.decode('utf-8') if sys.version_info < (3,) else query,
+               server_id=self.manager.sid,
-                                   query_id=query_id
+               conn_id=self.conn_id,
-                               )
+               query=query.decode('utf-8') if
-                               )
+               sys.version_info < (3,) else query,
+               query_id=query_id
+           )
+        )
         try:
             self.__internal_blocking_execute(cur, query, params)
         except psycopg2.Error as pe:
             cur.close()
             errmsg = self._formatted_exception_msg(pe, formatted_exception_msg)
             current_app.logger.error(
-                u"failed to execute query ((with server cursor) for the server #{server_id} - {conn_id} (query-id: {query_id}):\nerror message:{errmsg}".format(
+                u"failed to execute query ((with server cursor) "
+                u"for the server #{server_id} - {conn_id} "
+                u"(query-id: {query_id}):\nerror message:{errmsg}".format(
                     server_id=self.manager.sid,
                     conn_id=self.conn_id,
                     query=query,
@@ -632,6 +656,33 @@ WHERE
             )
             return False, errmsg
 
+        def handle_json_data(json_columns, results):
+            """
+            [ This is only for Python2.x]
+            This function will be useful to handle json data types.
+            We will dump json data as proper json instead of unicode values
+
+            Args:
+                json_columns: Columns which contains json data
+                results: Query result
+
+            Returns:
+                results
+            """
+            # Only if Python2 and there are columns with JSON type
+            if sys.version_info < (3,) and len(json_columns) > 0:
+                temp_results = []
+                for row in results:
+                    res = dict()
+                    for k, v in row.items():
+                        if k in json_columns:
+                            res[k] = json.dumps(v)
+                        else:
+                            res[k] = v
+                    temp_results.append(res)
+                results = temp_results
+            return results
+
         def gen():
 
             results = cur.fetchmany(records)
@@ -640,15 +691,26 @@ WHERE
                     cur.close()
                 yield gettext('The query executed did not return any data.')
                 return
-
+            header = []
-            header = [c.to_dict()['name'] for c in cur.ordered_description()]
+            json_columns = []
-
+            # json, jsonb, json[], jsonb[]
+            json_types = (114, 199, 3802, 3807)
+            for c in cur.ordered_description():
+                # This is to handle the case in which column name is non-ascii
+                header.append(u"" + c.to_dict()['name'])
+                if c.to_dict()['type_code'] in json_types:
+                    json_columns.append(
+                        u"" + c.to_dict()['name']
+                    )
             res_io = StringIO()
 
             csv_writer = csv.DictWriter(
-                res_io, fieldnames=header, delimiter=str(','), quoting=csv.QUOTE_NONNUMERIC
+                res_io, fieldnames=header, delimiter=u',',
+                quoting=csv.QUOTE_NONNUMERIC
             )
+
             csv_writer.writeheader()
+            results = handle_json_data(json_columns, results)
             csv_writer.writerows(results)
 
             yield res_io.getvalue()
@@ -663,8 +725,10 @@ WHERE
                 res_io = StringIO()
 
                 csv_writer = csv.DictWriter(
-                    res_io, fieldnames=header, delimiter=str(','), quoting=csv.QUOTE_NONNUMERIC
+                    res_io, fieldnames=header, delimiter=u',',
+                    quoting=csv.QUOTE_NONNUMERIC
                 )
+                results = handle_json_data(json_columns, results)
                 csv_writer.writerows(results)
                 yield res_io.getvalue()