import json
import peewee
from playhouse.migrate import Migrator
from redash import db
from redash import models

if __name__ == '__main__':
    default_options = {"series": {"type": "bar"}}

    db.connect_db()

    if not models.Visualization.table_exists():
        print "Creating visualization table..."
        models.Visualization.create_table()

    with db.database.transaction():
        migrator = Migrator(db.database)
        print "Adding visualization_id to widgets:"
        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,
                                   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"):
        vis = models.Visualization(query=query, name="Cohort",
                                   description=query.description,
                                   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()
        if vis:
            print '... visualization type (%s) found.' % vis_type
            widget.visualization = vis[0]
            widget.save()
        else:
            vis_name = vis_type.title()

            options = json.loads(widget.options)
            vis_options = {"series": options} if options else default_options
            vis_options = json.dumps(vis_options)

            vis = models.Visualization(query=query, name=vis_name,
                                       description=query.description,
                                       type=vis_type, options=vis_options)

            print '... Created visualization for type: %s' % vis_type
            vis.save()
            widget.visualization = vis
            widget.save()

        db.close_db(None)
@@ -9,22 +9,8 @@
 son%0A
-import peewee%0A
 from
@@ -462,16 +462,86 @@
 dgets:%22%0A
+        field = models.Widget.visualization%0A        field.null = True%0A
         
@@ -844,32 +844,38 @@
 uery.description
+ or %22%22
 ,%0A              
@@ -1074,28 +1074,8 @@
 ere(
-peewee.fn.Icontains(
 mode
@@ -1091,18 +1091,21 @@
 name
-)==%22
+ ** %22%25
 cohort
+%25
 %22):%0A
@@ -1223,32 +1223,38 @@
 uery.description
+ or %22%22
 ,%0A              
@@ -1497,37 +1497,8 @@
 .id%0A
-        query = widget.query%0A
     
@@ -1595,16 +1595,72 @@
 TABLE'%0A%0A
+        query = models.Query.get_by_id(widget.query_id)%0A
         
@@ -1692,21 +1692,45 @@
 s.where(
-type=
+models.Visualization.type == 
 vis_type
@@ -1863,11 +1863,8 @@
  vis
-%5B0%5D
 %0A   
@@ -2254,16 +2254,22 @@
 cription
+ or %22%22
 ,%0A      
@@ -2497,24 +2497,235 @@
 save()%0A%0A    
+with db.database.transaction():%0A        migrator = Migrator(db.database)%0A        print %22Setting visualization_id as not null...%22%0A        migrator.set_nullable(models.Widget, models.Widget.visualization, False)%0A%0A
     db.close

Improve visualizations migration