Subject: [pgAdmin4][PATCH] To fix the issues with CSV file download

From: Murtuza Zabuawala <[email-masked]>
Date: Mon, 8 May 2017 17:44:00 +0530

Hi,

Please find the patch to fix the issues in CSV file download feature,

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

2) To handle non-ascii query data
RM#2353

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

Please review.

--
Regards,
Murtuza Zabuawala
EnterpriseDB: http://www.enterprisedb.com
The Enterprise PostgreSQL Company
diff --git a/requirements.txt b/requirements.txt
index 8afeb52..90ed057 100644
--- a/requirements.txt
+++ b/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'
\ No newline at end of file
diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py
index d114988..2d1d4e5 100644
--- a/web/pgadmin/tools/sqleditor/__init__.py
+++ b/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.")
+        )
diff --git a/web/pgadmin/utils/driver/psycopg2/__init__.py b/web/pgadmin/utils/driver/psycopg2/__init__.py
index e474817..f2e4b72 100644
--- a/web/pgadmin/utils/driver/psycopg2/__init__.py
+++ b/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
 
@@ -596,7 +600,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
 
@@ -607,21 +626,26 @@ WHERE
         if sys.version_info < (3,) and type(query) == unicode:
             query = query.encode('utf-8')
 
-        current_app.logger.log(25,
-                               u"Execute (with server cursor) for server #{server_id} - {conn_id} (Query-id: {query_id}):\n{query}".format(
-                                   server_id=self.manager.sid,
-                                   conn_id=self.conn_id,
-                                   query=query.decode('utf-8') if sys.version_info < (3,) else query,
-                                   query_id=query_id
-                               )
-                               )
+        current_app.logger.log(
+           25,
+           u"Execute (with server cursor) for server #{server_id} - {conn_id} "
+           u"(Query-id: {query_id}):\n{query}".format(
+               server_id=self.manager.sid,
+               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,
@@ -631,6 +655,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)
@@ -639,15 +690,26 @@ WHERE
                     cur.close()
                 yield gettext('The query executed did not return any data.')
                 return
-
-            header = [c.to_dict()['name'] for c in cur.ordered_description()]
-
+            header = []
+            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()
@@ -662,8 +724,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()
 
@@ -1487,8 +1551,8 @@ class ServerManager(object):
                     not isinstance(database, unicode):
                 database = database.decode('utf-8')
             if did is not None:
-                if did in self.db_info:                        
-                    self.db_info[did]['datname']=database            
+                if did in self.db_info:
+                    self.db_info[did]['datname']=database
         else:
             if did is None:
                 database = self.db

-- 
Sent via pgadmin-hackers mailing list ([email-masked])
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgadmin-hackers

---

From: Dave Page <[email-masked]>
Date: Mon, 8 May 2017 13:36:30 +0100

Thanks, patch applied.

-- 
Dave Page
Blog: http://pgsnake.blogspot.com
Twitter: @pgsnake

EnterpriseDB UK: http://www.enterprisedb.com
The Enterprise PostgreSQL Company

---

From: Murtuza Zabuawala <[email-masked]>
Date: Thu, 11 May 2017 19:51:34 +0530

Hi Dave,

Please find patch to handle non-ascii column names in CSV for Pyhton2.

Thanks to Harshal for pointing out the issue.

On Mon, May 8, 2017 at 6:06 PM, Dave Page <[email-masked]> wrote:

> Thanks, patch applied.
>
> On Mon, May 8, 2017 at 1:14 PM, Murtuza Zabuawala <murtuza.zabuawala@
> enterprisedb.com> wrote:
>
>> Hi,
>>
diff --git a/web/pgadmin/utils/driver/psycopg2/__init__.py b/web/pgadmin/utils/driver/psycopg2/__init__.py
index 1f7ccce..502cee4 100644
--- a/web/pgadmin/utils/driver/psycopg2/__init__.py
+++ b/web/pgadmin/utils/driver/psycopg2/__init__.py
@@ -42,9 +42,11 @@ if sys.version_info < (3,):
     from StringIO import StringIO
     psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
     psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
+    IS_PY2 = True
 else:
     from io import StringIO
     import csv
+    IS_PY2 = False
 
 _ = gettext
 
@@ -624,7 +626,7 @@ WHERE
             return False, str(cur)
         query_id = random.randint(1, 9999999)
 
-        if sys.version_info < (3,) and type(query) == unicode:
+        if IS_PY2 and type(query) == unicode:
             query = query.encode('utf-8')
 
         current_app.logger.log(
@@ -670,7 +672,7 @@ WHERE
                 results
             """
             # Only if Python2 and there are columns with JSON type
-            if sys.version_info < (3,) and len(json_columns) > 0:
+            if IS_PY2 and len(json_columns) > 0:
                 temp_results = []
                 for row in results:
                     res = dict()
@@ -683,6 +685,26 @@ WHERE
                 results = temp_results
             return results
 
+        def convert_keys_to_unicode(results, conn_encoding):
+            """
+            [ This is only for Python2.x]
+            We need to convert all keys to unicode as psycopg2
+            sends them as string
+
+            Args:
+                res: Query result set from psycopg2
+                conn_encoding: Connection encoding
+
+            Returns:
+                Result set (With all the keys converted to unicode)
+            """
+            new_results = []
+            for row in results:
+                new_results.append(
+                    dict([(k.decode(conn_encoding), v) for k, v in row.items()])
+                )
+            return new_results
+
         def gen():
 
             results = cur.fetchmany(records)
@@ -691,17 +713,25 @@ WHERE
                     cur.close()
                 yield gettext('The query executed did not return any data.')
                 return
+
             header = []
             json_columns = []
+            conn_encoding = cur.connection.encoding
+
             # 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'])
+                column_name = c.to_dict()['name']
+                if IS_PY2:
+                    column_name = column_name.decode(conn_encoding)
+                header.append(column_name)
                 if c.to_dict()['type_code'] in json_types:
-                    json_columns.append(
-                        u"" + c.to_dict()['name']
-                    )
+                    json_columns.append(column_name)
+
+            if IS_PY2:
+                results = convert_keys_to_unicode(results, conn_encoding)
+
             res_io = StringIO()
 
             csv_writer = csv.DictWriter(
@@ -728,6 +758,10 @@ WHERE
                     res_io, fieldnames=header, delimiter=u',',
                     quoting=csv.QUOTE_NONNUMERIC
                 )
+
+                if IS_PY2:
+                    results = convert_keys_to_unicode(results, conn_encoding)
+
                 results = handle_json_data(json_columns, results)
                 csv_writer.writerows(results)
                 yield res_io.getvalue()

---

From: Dave Page <[email-masked]>
Date: Fri, 12 May 2017 10:38:32 +0100

Thanks - patch applied.