Merge pull request #131 from EverythingMe/feature_import

Feature: import dashboard (along with widgets, visualization and queries) from JSON

manage.py

@@ -10,6 +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 import_manager
 from flask.ext.script import Manager, prompt_pass
 
 manager = Manager(app)
@@ -93,6 +94,7 @@ def delete(email):
 
 manager.add_command("database", database_manager)
 manager.add_command("users", users_manager)
+manager.add_command("import", import_manager)
 
 if __name__ == '__main__':
     channel = logging.StreamHandler()

redash/import_export.py

@@ -0,0 +1,125 @@
+import json
+from redash import models
+from flask.ext.script import Manager
+
+
+class Importer(object):
+    def __init__(self, object_mapping=None):
+        if object_mapping is None:
+            object_mapping = {}
+        self.object_mapping = object_mapping
+
+    def import_query_result(self, query_result):
+        query_result = self._get_or_create(models.QueryResult, query_result['id'],
+                                           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(self, user, query):
+        query_result = self.import_query_result(query['latest_query_data'])
+
+        new_query = self._get_or_create(models.Query, query['id'], 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
+
+
+    def import_visualization(self, user, visualization):
+        query = self.import_query(user, visualization['query'])
+
+        new_visualization = self._get_or_create(models.Visualization, visualization['id'],
+                                                name=visualization['name'],
+                                                description=visualization['description'],
+                                                type=visualization['type'],
+                                                options=json.dumps(visualization['options']),
+                                                query=query)
+        return new_visualization
+
+    def import_widget(self, dashboard, widget):
+        visualization = self.import_visualization(dashboard.user, widget['visualization'])
+
+        new_widget = self._get_or_create(models.Widget, widget['id'],
+                                         dashboard=dashboard,
+                                         width=widget['width'],
+                                         options=json.dumps(widget['options']),
+                                         visualization=visualization)
+
+        return new_widget
+
+    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).
+        """
+
+        new_dashboard = self._get_or_create(models.Dashboard, dashboard['id'],
+                                            name=dashboard['name'],
+                                            slug=dashboard['slug'],
+                                            layout='[]',
+                                            user=user)
+
+        layout = []
+
+        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)
+
+        new_dashboard.layout = json.dumps(layout)
+        new_dashboard.save()
+
+        return new_dashboard
+
+    def _get_or_create(self, object_type, external_id, **properties):
+        internal_id = self._get_mapping(object_type, external_id)
+        if internal_id:
+            update = object_type.update(**properties).where(object_type.id == internal_id)
+            update.execute()
+            obj = object_type.get_by_id(internal_id)
+        else:
+            obj = object_type.create(**properties)
+            self._update_mapping(object_type, external_id, obj.id)
+
+        return obj
+
+    def _get_mapping(self, object_type, external_id):
+        self.object_mapping.setdefault(object_type.__name__, {})
+        return self.object_mapping[object_type.__name__].get(str(external_id), None)
+
+    def _update_mapping(self, object_type, external_id, internal_id):
+        self.object_mapping.setdefault(object_type.__name__, {})
+        self.object_mapping[object_type.__name__][str(external_id)] = internal_id
+
+import_manager = Manager(help="import utilities")
+export_manager = Manager(help="export utilities")
+
+@import_manager.command
+def dashboard(mapping_filename, dashboard_filename, user_id):
+    with open(mapping_filename) as f:
+        mapping = json.loads(f.read())
+
+    user = models.User.get_by_id(user_id)
+    with open(dashboard_filename) as f:
+        dashboard = json.loads(f.read())
+
+        importer = Importer(object_mapping=mapping)
+        importer.import_dashboard(user, dashboard)
+
+    with open(mapping_filename, 'w') as f:
+        f.write(json.dumps(importer.object_mapping, indent=2))
+

redash/models.py

@@ -68,7 +68,7 @@ class ActivityLog(BaseModel):
         return unicode(self.id)
 
 
-class QueryResult(db.Model):
+class QueryResult(BaseModel):
     id = peewee.PrimaryKeyField()
     query_hash = peewee.CharField(max_length=32, index=True)
     query = peewee.TextField()
@@ -271,7 +271,7 @@ class Visualization(BaseModel):
         return u"%s %s" % (self.id, self.type)
 
 
-class Widget(db.Model):
+class Widget(BaseModel):
     id = peewee.PrimaryKeyField()
     visualization = peewee.ForeignKeyField(Visualization, related_name='widgets')
 

tests/test_import.py

@@ -0,0 +1,53 @@
+import json
+import os.path
+from tests import BaseTestCase
+from redash import models
+from redash import import_export
+from factories import user_factory, dashboard_factory
+
+
+class ImportTest(BaseTestCase):
+    def setUp(self):
+        super(ImportTest, self).setUp()
+
+        with open(os.path.join(os.path.dirname(__file__), 'flights.json')) as f:
+            self.dashboard = json.loads(f.read())
+            self.user = user_factory.create()
+
+    def test_imports_dashboard_correctly(self):
+        importer = import_export.Importer()
+        dashboard = importer.import_dashboard(self.user, self.dashboard)
+
+        self.assertIsNotNone(dashboard)
+        self.assertEqual(dashboard.name, self.dashboard['name'])
+        self.assertEqual(dashboard.slug, self.dashboard['slug'])
+        self.assertEqual(dashboard.user, self.user)
+
+        self.assertEqual(dashboard.widgets.count(),
+                         reduce(lambda s, row: s + len(row), self.dashboard['widgets'], 0))
+
+        self.assertEqual(models.Visualization.select().count(), dashboard.widgets.count())
+        self.assertEqual(models.Query.select().count(), dashboard.widgets.count()-1)
+        self.assertEqual(models.QueryResult.select().count(), dashboard.widgets.count()-1)
+
+    def test_imports_updates_existing_models(self):
+        importer = import_export.Importer()
+        importer.import_dashboard(self.user, self.dashboard)
+
+        self.dashboard['name'] = 'Testing #2'
+        dashboard = importer.import_dashboard(self.user, self.dashboard)
+        self.assertEqual(dashboard.name, self.dashboard['name'])
+        self.assertEquals(models.Dashboard.select().count(), 1)
+
+    def test_using_existing_mapping(self):
+        dashboard = dashboard_factory.create()
+        mapping = {
+            'Dashboard': {
+                "1": dashboard.id
+            }
+        }
+
+        importer = import_export.Importer(object_mapping=mapping)
+        imported_dashboard = importer.import_dashboard(self.user, self.dashboard)
+
+        self.assertEqual(imported_dashboard, dashboard)