From dfed319ac6fcfcddc04e63ddb8f526b5d8564158 Mon Sep 17 00:00:00 2001 From: JulieRossi Date: Fri, 8 Sep 2017 16:06:04 +0200 Subject: [PATCH] config: allow changing default options for writting csv --- superset/config.py | 5 +++++ superset/views/core.py | 4 ++-- superset/viz.py | 2 +- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/superset/config.py b/superset/config.py index 94d8dfc56a673..22e0d240558de 100644 --- a/superset/config.py +++ b/superset/config.py @@ -172,6 +172,11 @@ ENABLE_CORS = False CORS_OPTIONS = {} +# CSV Options: key/value pairs that will be passed as argument to DataFrame.to_csv method +# note: index option should not be overridden +CSV_EXPORT = { + 'encoding': 'utf-8', +} # --------------------------------------------------- # List of viz_types not allowed in your environment diff --git a/superset/views/core.py b/superset/views/core.py index 68e82027d0658..72d2cf6cafbb5 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2126,13 +2126,13 @@ def csv(self, client_id): columns = [c['name'] for c in obj['columns']] df = pd.DataFrame.from_records(obj['data'], columns=columns) logging.info("Using pandas to convert to CSV") - csv = df.to_csv(index=False, encoding='utf-8') + csv = df.to_csv(index=False, **config.get('CSV_EXPORT')) else: logging.info("Running a query to turn into CSV") sql = query.select_sql or query.executed_sql df = query.database.get_df(sql, query.schema) # TODO(bkyryliuk): add compression=gzip for big files. - csv = df.to_csv(index=False, encoding='utf-8') + csv = df.to_csv(index=False, **config.get('CSV_EXPORT')) response = Response(csv, mimetype='text/csv') response.headers['Content-Disposition'] = ( 'attachment; filename={}.csv'.format(query.name)) diff --git a/superset/viz.py b/superset/viz.py index 7a2b22ac7f2e7..57606e81d02fb 100644 --- a/superset/viz.py +++ b/superset/viz.py @@ -310,7 +310,7 @@ def data(self): def get_csv(self): df = self.get_df() include_index = not isinstance(df.index, pd.RangeIndex) - return df.to_csv(index=include_index, encoding="utf-8") + return df.to_csv(index=include_index, **config.get('CSV_EXPORT')) def get_data(self, df): return []