From d6bd19438c6ab8b856797cf1733cccf2a4fe911a Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Tue, 11 Mar 2014 15:13:41 +0200 Subject: [PATCH] Move import functions into a class, to have state --- manage.py | 4 +- redash/import_export.py | 125 ++++++++++++++++++++++------------------ tests/test_import.py | 3 +- 3 files changed, 74 insertions(+), 58 deletions(-) diff --git a/manage.py b/manage.py index dc057a0e..bb5f4eeb 100755 --- a/manage.py +++ b/manage.py @@ -10,7 +10,7 @@ atfork.stdlib_fixer.fix_logging_module() import logging import time from redash import settings, app, db, models, data_manager, __version__ -from redash.import_export import manager as import_export_manager +from redash.import_export import import_manager from flask.ext.script import Manager, prompt_pass manager = Manager(app) @@ -94,7 +94,7 @@ def delete(email): manager.add_command("database", database_manager) manager.add_command("users", users_manager) -manager.add_command("import", import_export_manager) +manager.add_command("import", import_manager) if __name__ == '__main__': channel = logging.StreamHandler() diff --git a/redash/import_export.py b/redash/import_export.py index 299ceef2..61a5ab0d 100644 --- a/redash/import_export.py +++ b/redash/import_export.py @@ -3,86 +3,101 @@ from redash import models from flask.ext.script import Manager -def import_query_result(query_result): - query_result = models.QueryResult.create(data=json.dumps(query_result['data']), - query_hash=query_result['query_hash'], - retrieved_at=query_result['retrieved_at'], - query=query_result['query'], - runtime=query_result['runtime']) +class Importer(object): + def __init__(self, object_mapping={}): + self.object_mapping = object_mapping - return query_result + def import_query_result(self, query_result): + query_result = models.QueryResult.create(data=json.dumps(query_result['data']), + query_hash=query_result['query_hash'], + retrieved_at=query_result['retrieved_at'], + query=query_result['query'], + runtime=query_result['runtime']) + + return query_result -def import_query(user, query): - query_result = import_query_result(query['latest_query_data']) + def import_query(self, user, query): + query_result = self.import_query_result(query['latest_query_data']) - new_query = models.Query.create(name=query['name'], - user=user, - ttl=-1, - query=query['query'], - query_hash=query['query_hash'], - description=query['description'], - latest_query_data=query_result) + new_query = models.Query.create(name=query['name'], + user=user, + ttl=-1, + query=query['query'], + query_hash=query['query_hash'], + description=query['description'], + latest_query_data=query_result) - return new_query + return new_query -def import_visualization(user, visualization): - query = import_query(user, visualization['query']) + def import_visualization(self, user, visualization): + query = self.import_query(user, visualization['query']) - new_visualization = models.Visualization.create(name=visualization['name'], - description=visualization['description'], - type=visualization['type'], - options=json.dumps(visualization['options']), - query=query) - return new_visualization + new_visualization = models.Visualization.create(name=visualization['name'], + description=visualization['description'], + type=visualization['type'], + options=json.dumps(visualization['options']), + query=query) + return new_visualization -def import_widget(dashboard, widget): - visualization = import_visualization(dashboard.user, widget['visualization']) + def import_widget(self, dashboard, widget): + visualization = self.import_visualization(dashboard.user, widget['visualization']) - new_widget = models.Widget.create(dashboard=dashboard, - width=widget['width'], - options=json.dumps(widget['options']), - visualization=visualization) + new_widget = models.Widget.create(dashboard=dashboard, + width=widget['width'], + options=json.dumps(widget['options']), + visualization=visualization) - return new_widget + return new_widget -def import_dashboard(user, dashboard): - """ - Imports dashboard along with widgets, visualizations and queries from another re:dash. + def import_dashboard(self, user, dashboard): + """ + Imports dashboard along with widgets, visualizations and queries from another re:dash. - user - the user to associate all objects with. - dashboard - dashboard to import (can be result of loading a json output). - """ + user - the user to associate all objects with. + dashboard - dashboard to import (can be result of loading a json output). + """ - new_dashboard = models.Dashboard.create(name=dashboard['name'], - slug=dashboard['slug'], - layout='[]', - user=user) + new_dashboard = models.Dashboard.create(name=dashboard['name'], + slug=dashboard['slug'], + layout='[]', + user=user) - layout = [] + layout = [] - for widgets in dashboard['widgets']: - row = [] - for widget in widgets: - widget_id = import_widget(new_dashboard, widget).id - row.append(widget_id) + for widgets in dashboard['widgets']: + row = [] + for widget in widgets: + widget_id = self.import_widget(new_dashboard, widget).id + row.append(widget_id) - layout.append(row) + layout.append(row) - new_dashboard.layout = json.dumps(layout) - new_dashboard.save() + new_dashboard.layout = json.dumps(layout) + new_dashboard.save() - return new_dashboard + return new_dashboard -manager = Manager(help="import/export utilities") + def _get_mapping(self, object_type, external_id): + self.object_mapping.setdefault(object_type, {}) + return self.object_mapping[object_type].get(external_id, None) -@manager.command + def _update_mapping(self, object_type, external_id, internal_id): + self.object_mapping.setdefault(object_type, {}) + self.object_mapping[object_type][external_id] = internal_id + +import_manager = Manager(help="import utilities") +export_manager = Manager(help="export utilities") + +@import_manager.command def dashboard(filename, user_id): user = models.User.get_by_id(user_id) with open(filename) as f: dashboard = json.loads(f.read()) - import_dashboard(user, dashboard) \ No newline at end of file + importer = Importer() + importer.import_dashboard(user, dashboard) + diff --git a/tests/test_import.py b/tests/test_import.py index 7e26d6f4..a3837739 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -14,7 +14,8 @@ class ImportTest(BaseTestCase): self.user = user_factory.create() def test_imports_dashboard_correctly(self): - dashboard = import_export.import_dashboard(self.user, self.dashboard) + importer = import_export.Importer() + dashboard = importer.import_dashboard(self.user, self.dashboard) self.assertIsNotNone(dashboard) self.assertEqual(dashboard.name, self.dashboard['name'])