pgadmin-org
/
pgadmin4
mirror of https://github.com/pgadmin-org/pgadmin4.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
pull/5/head
Murtuza Zabuawala 2017-05-08 13:36:11 +01:00 committed by Dave Page
parent a80f760933
commit 63d42745ef
3 changed files with 106 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
requirements.txt
View File

@ -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'

28
web/pgadmin/tools/sqleditor/__init__.py
View File

@ -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.")
)

100
web/pgadmin/utils/driver/psycopg2/__init__.py
View File

@ -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,
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,
@ -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 = [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()
@ -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()
@ -1488,8 +1552,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