From 12c30ed487421dc46ceb512369028b27d116929f Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Thu, 18 Jan 2018 18:28:02 -0800 Subject: [PATCH 1/2] make superset imports and exports use json, not pickle --- superset/models/core.py | 20 ++++++++++++++++---- superset/utils.py | 38 ++++++++++++++++++++++++++++++++++++++ superset/views/core.py | 17 ++++++++++------- 3 files changed, 64 insertions(+), 11 deletions(-) diff --git a/superset/models/core.py b/superset/models/core.py index 413f85bf99c81..397235480bd58 100644 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -9,7 +9,6 @@ import functools import json import logging -import pickle import textwrap from flask import escape, g, Markup, request @@ -395,7 +394,7 @@ def import_obj(cls, dashboard_to_import, import_time=None): be overridden or just copies over. Slices that belong to this dashboard will be wired to existing tables. This function can be used to import/export dashboards between multiple superset instances. - Audit metadata isn't copies over. + Audit metadata isn't copied over. """ def alter_positions(dashboard, old_to_new_slc_id_dict): """ Updates slice_ids in the position json. @@ -499,6 +498,19 @@ def alter_positions(dashboard, old_to_new_slc_id_dict): @classmethod def export_dashboards(cls, dashboard_ids): + + class DashboardEncoder(json.JSONEncoder): + # pylint: disable=E0202 + def default(self, o): + try: + vals = { + k: v for k, v in o.__dict__.items() if k != '_sa_instance_state'} + return {'__{}__'.format(o.__class__.__name__): vals} + except Exception as e: + if type(o) == datetime: + return {'__datetime__': o.replace(microsecond=0).isoformat()} + return json.JSONEncoder.default(self, o) + copied_dashboards = [] datasource_ids = set() for dashboard_id in dashboard_ids: @@ -533,10 +545,10 @@ def export_dashboards(cls, dashboard_ids): make_transient(eager_datasource) eager_datasources.append(eager_datasource) - return pickle.dumps({ + return json.dumps({ 'dashboards': copied_dashboards, 'datasources': eager_datasources, - }) + }, cls=DashboardEncoder, indent=4) class Database(Model, AuditMixinNullable, ImportMixin): diff --git a/superset/utils.py b/superset/utils.py index 8224843213d2a..928e2f01fed98 100644 --- a/superset/utils.py +++ b/superset/utils.py @@ -42,6 +42,8 @@ from sqlalchemy import event, exc, select from sqlalchemy.types import TEXT, TypeDecorator + + logging.getLogger('MARKDOWN').setLevel(logging.INFO) PY3K = sys.version_info >= (3, 0) @@ -240,6 +242,42 @@ def dttm_from_timtuple(d): d.tm_year, d.tm_mon, d.tm_mday, d.tm_hour, d.tm_min, d.tm_sec) +def decode_dashboards(o): + """ + Function to be passed into json.loads obj_hook parameter + Recreates the dashboard object from a json representation. + """ + import superset.models.core as models + from superset.connectors.sqla.models import ( + SqlaTable, SqlMetric, TableColumn, + ) + + if '__Dashboard__' in o: + d = models.Dashboard() + d.__dict__.update(o['__Dashboard__']) + return d + elif '__Slice__' in o: + d = models.Slice() + d.__dict__.update(o['__Slice__']) + return d + elif '__TableColumn__' in o: + d = TableColumn() + d.__dict__.update(o['__TableColumn__']) + return d + elif '__SqlaTable__' in o: + d = SqlaTable() + d.__dict__.update(o['__SqlaTable__']) + return d + elif '__SqlMetric__' in o: + d = SqlMetric() + d.__dict__.update(o['__SqlMetric__']) + return d + elif '__datetime__' in o: + return datetime.strptime(o['__datetime__'], '%Y-%m-%dT%H:%M:%S') + else: + return o + + def parse_human_timedelta(s): """ Returns ``datetime.datetime`` from natural language time deltas diff --git a/superset/views/core.py b/superset/views/core.py index ec4cce1fb1568..d372992281268 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -8,7 +8,6 @@ import json import logging import os -import pickle import re import time import traceback @@ -38,7 +37,9 @@ viz, ) from superset.connectors.connector_registry import ConnectorRegistry -from superset.connectors.sqla.models import AnnotationDatasource, SqlaTable +from superset.connectors.sqla.models import ( + AnnotationDatasource, SqlaTable, SqlMetric, TableColumn, +) from superset.forms import CsvToDatabaseForm from superset.legacy import cast_form_data import superset.models.core as models @@ -601,7 +602,7 @@ def download_dashboards(self): ids = request.args.getlist('id') return Response( models.Dashboard.export_dashboards(ids), - headers=generate_download_headers('pickle'), + headers=generate_download_headers('json'), mimetype='application/text') return self.render_template( 'superset/export_dashboards.html', @@ -1114,15 +1115,17 @@ def explore_json(self, datasource_type, datasource_id): @has_access @expose('/import_dashboards', methods=['GET', 'POST']) def import_dashboards(self): - """Overrides the dashboards using pickled instances from the file.""" + """Overrides the dashboards using json instances from the file.""" + + + f = request.files.get('file') if request.method == 'POST' and f: current_tt = int(time.time()) - data = pickle.load(f) + data = json.loads(f.stream.read(), object_hook=utils.decode_dashboards) # TODO: import DRUID datasources for table in data['datasources']: - ds_class = ConnectorRegistry.sources.get(table.type) - ds_class.import_obj(table, import_time=current_tt) + type(table).import_obj(table, import_time=current_tt) db.session.commit() for dashboard in data['dashboards']: models.Dashboard.import_obj( From 48a79e176c18c6f7caabddb03ae73d44a3dd26c6 Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Mon, 22 Jan 2018 00:12:42 -0800 Subject: [PATCH 2/2] fix tests --- superset/models/core.py | 15 +-------------- superset/utils.py | 14 +++++++++++++- superset/views/core.py | 7 +------ tests/import_export_tests.py | 33 ++++++++++++++++++++++++--------- 4 files changed, 39 insertions(+), 30 deletions(-) diff --git a/superset/models/core.py b/superset/models/core.py index 397235480bd58..c55cbf8977b89 100644 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -498,19 +498,6 @@ def alter_positions(dashboard, old_to_new_slc_id_dict): @classmethod def export_dashboards(cls, dashboard_ids): - - class DashboardEncoder(json.JSONEncoder): - # pylint: disable=E0202 - def default(self, o): - try: - vals = { - k: v for k, v in o.__dict__.items() if k != '_sa_instance_state'} - return {'__{}__'.format(o.__class__.__name__): vals} - except Exception as e: - if type(o) == datetime: - return {'__datetime__': o.replace(microsecond=0).isoformat()} - return json.JSONEncoder.default(self, o) - copied_dashboards = [] datasource_ids = set() for dashboard_id in dashboard_ids: @@ -548,7 +535,7 @@ def default(self, o): return json.dumps({ 'dashboards': copied_dashboards, 'datasources': eager_datasources, - }, cls=DashboardEncoder, indent=4) + }, cls=utils.DashboardEncoder, indent=4) class Database(Model, AuditMixinNullable, ImportMixin): diff --git a/superset/utils.py b/superset/utils.py index 928e2f01fed98..e28eda30b5d28 100644 --- a/superset/utils.py +++ b/superset/utils.py @@ -43,7 +43,6 @@ from sqlalchemy.types import TEXT, TypeDecorator - logging.getLogger('MARKDOWN').setLevel(logging.INFO) PY3K = sys.version_info >= (3, 0) @@ -278,6 +277,19 @@ def decode_dashboards(o): return o +class DashboardEncoder(json.JSONEncoder): + # pylint: disable=E0202 + def default(self, o): + try: + vals = { + k: v for k, v in o.__dict__.items() if k != '_sa_instance_state'} + return {'__{}__'.format(o.__class__.__name__): vals} + except Exception: + if type(o) == datetime: + return {'__datetime__': o.replace(microsecond=0).isoformat()} + return json.JSONEncoder.default(self, o) + + def parse_human_timedelta(s): """ Returns ``datetime.datetime`` from natural language time deltas diff --git a/superset/views/core.py b/superset/views/core.py index d372992281268..b06492cb626c5 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -37,9 +37,7 @@ viz, ) from superset.connectors.connector_registry import ConnectorRegistry -from superset.connectors.sqla.models import ( - AnnotationDatasource, SqlaTable, SqlMetric, TableColumn, -) +from superset.connectors.sqla.models import AnnotationDatasource, SqlaTable from superset.forms import CsvToDatabaseForm from superset.legacy import cast_form_data import superset.models.core as models @@ -1116,9 +1114,6 @@ def explore_json(self, datasource_type, datasource_id): @expose('/import_dashboards', methods=['GET', 'POST']) def import_dashboards(self): """Overrides the dashboards using json instances from the file.""" - - - f = request.files.get('file') if request.method == 'POST' and f: current_tt = int(time.time()) diff --git a/tests/import_export_tests.py b/tests/import_export_tests.py index d51b9592976df..245d4199908fa 100644 --- a/tests/import_export_tests.py +++ b/tests/import_export_tests.py @@ -5,12 +5,11 @@ from __future__ import unicode_literals import json -import pickle import unittest from sqlalchemy.orm.session import make_transient -from superset import db +from superset import db, utils from superset.connectors.druid.models import ( DruidColumn, DruidDatasource, DruidMetric, ) @@ -205,13 +204,22 @@ def test_export_1_dashboard(self): .format(birth_dash.id) ) resp = self.client.get(export_dash_url) - exported_dashboards = pickle.loads(resp.data)['dashboards'] + exported_dashboards = json.loads( + resp.data.decode('utf-8'), + object_hook=utils.decode_dashboards, + )['dashboards'] self.assert_dash_equals(birth_dash, exported_dashboards[0]) self.assertEquals( birth_dash.id, - json.loads(exported_dashboards[0].json_metadata)['remote_id']) - - exported_tables = pickle.loads(resp.data)['datasources'] + json.loads( + exported_dashboards[0].json_metadata, + object_hook=utils.decode_dashboards, + )['remote_id']) + + exported_tables = json.loads( + resp.data.decode('utf-8'), + object_hook=utils.decode_dashboards, + )['datasources'] self.assertEquals(1, len(exported_tables)) self.assert_table_equals( self.get_table_by_name('birth_names'), exported_tables[0]) @@ -223,8 +231,12 @@ def test_export_2_dashboards(self): '/dashboardmodelview/export_dashboards_form?id={}&id={}&action=go' .format(birth_dash.id, world_health_dash.id)) resp = self.client.get(export_dash_url) - exported_dashboards = sorted(pickle.loads(resp.data)['dashboards'], - key=lambda d: d.dashboard_title) + exported_dashboards = sorted( + json.loads( + resp.data.decode('utf-8'), + object_hook=utils.decode_dashboards, + )['dashboards'], + key=lambda d: d.dashboard_title) self.assertEquals(2, len(exported_dashboards)) self.assert_dash_equals(birth_dash, exported_dashboards[0]) self.assertEquals( @@ -239,7 +251,10 @@ def test_export_2_dashboards(self): ) exported_tables = sorted( - pickle.loads(resp.data)['datasources'], key=lambda t: t.table_name) + json.loads( + resp.data.decode('utf-8'), + object_hook=utils.decode_dashboards)['datasources'], + key=lambda t: t.table_name) self.assertEquals(2, len(exported_tables)) self.assert_table_equals( self.get_table_by_name('birth_names'), exported_tables[0])