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