From c40cf2e7e8601d9231beda3e196df7dbc262eb09 Mon Sep 17 00:00:00 2001
From: Arik Fraimovich <arik@arikfr.com>
Date: Sun, 9 Feb 2014 16:40:26 +0200
Subject: [PATCH] Improve visualizations migration

---
 migrations/create_visualizations.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/migrations/create_visualizations.py b/migrations/create_visualizations.py
index fc1847b42..3f743d75e 100644
--- a/migrations/create_visualizations.py
+++ b/migrations/create_visualizations.py
@@ -1,5 +1,4 @@
 import json
-import peewee
 from playhouse.migrate import Migrator
 from redash import db
 from redash import models
@@ -16,34 +15,36 @@ if __name__ == '__main__':
     with db.database.transaction():
         migrator = Migrator(db.database)
         print "Adding visualization_id to widgets:"
+        field = models.Widget.visualization
+        field.null = True
         migrator.add_column(models.Widget, models.Widget.visualization, 'visualization_id')
 
     print 'Creating TABLE visualizations for all queries...'
     for query in models.Query.select():
         vis = models.Visualization(query=query, name="Table",
-                                   description=query.description,
+                                   description=query.description or "",
                                    type="TABLE", options="{}")
         vis.save()
 
     print 'Creating COHORT visualizations for all queries named like %cohort%...'
-    for query in models.Query.select().where(peewee.fn.Icontains(models.Query.name)=="cohort"):
+    for query in models.Query.select().where(models.Query.name ** "%cohort%"):
         vis = models.Visualization(query=query, name="Cohort",
-                                   description=query.description,
+                                   description=query.description or "",
                                    type="COHORT", options="{}")
         vis.save()
 
     print 'Create visualization for all widgets (unless exists already):'
     for widget in models.Widget.select():
         print 'Processing widget id: %d:' % widget.id
-        query = widget.query
         vis_type = widget.type.upper()
         if vis_type == 'GRID':
             vis_type = 'TABLE'
 
-        vis = query.visualizations.where(type=vis_type).first()
+        query = models.Query.get_by_id(widget.query_id)
+        vis = query.visualizations.where(models.Visualization.type == vis_type).first()
         if vis:
             print '... visualization type (%s) found.' % vis_type
-            widget.visualization = vis[0]
+            widget.visualization = vis
             widget.save()
         else:
             vis_name = vis_type.title()
@@ -53,7 +54,7 @@ if __name__ == '__main__':
             vis_options = json.dumps(vis_options)
 
             vis = models.Visualization(query=query, name=vis_name,
-                                       description=query.description,
+                                       description=query.description or "",
                                        type=vis_type, options=vis_options)
 
             print '... Created visualization for type: %s' % vis_type
@@ -61,4 +62,9 @@ if __name__ == '__main__':
             widget.visualization = vis
             widget.save()
 
-        db.close_db(None)
\ No newline at end of file
+    with db.database.transaction():
+        migrator = Migrator(db.database)
+        print "Setting visualization_id as not null..."
+        migrator.set_nullable(models.Widget, models.Widget.visualization, False)
+
+    db.close_db(None)
\ No newline at end of file