[annotations] Fixing migration for annotation layers (#4187)

apache · Jan 10, 2018 · 0cb7c5e · 0cb7c5e
1 parent e182f7f
commit 0cb7c5e
Showing 1 changed file with 25 additions and 12 deletions.
diff --git a/superset/migrations/versions/21e88bc06c02_annotation_migration.py b/superset/migrations/versions/21e88bc06c02_annotation_migration.py
@@ -37,18 +37,19 @@ def upgrade():
             Slice.viz_type.like('line'), Slice.viz_type.like('bar'))):
         params = json.loads(slc.params)
         layers = params.get('annotation_layers', [])
-        new_layers = []
-        if len(layers) and isinstance(layers[0], int):
+        if layers:
+            new_layers = []
             for layer in layers:
-                new_layers.append(
-                    {
-                        'annotationType': 'INTERVAL',
-                        'style': 'solid',
-                        'name': 'Layer {}'.format(layer),
-                        'show': True,
-                        'overrides': {'since': None, 'until': None},
-                        'value': 1, 'width': 1, 'sourceType': 'NATIVE',
-                    })
+                new_layers.append({
+                    'annotationType': 'INTERVAL',
+                    'style': 'solid',
+                    'name': 'Layer {}'.format(layer),
+                    'show': True,
+                    'overrides': {'since': None, 'until': None},
+                    'value': layer,
+                    'width': 1,
+                    'sourceType': 'NATIVE',
+                })
             params['annotation_layers'] = new_layers
             slc.params = json.dumps(params)
             session.merge(slc)
@@ -57,4 +58,16 @@ def upgrade():
 
 
 def downgrade():
-    pass
+    bind = op.get_bind()
+    session = db.Session(bind=bind)
+
+    for slc in session.query(Slice).filter(or_(
+            Slice.viz_type.like('line'), Slice.viz_type.like('bar'))):
+        params = json.loads(slc.params)
+        layers = params.get('annotation_layers', [])
+        if layers:
+            params['annotation_layers'] = [layer['value'] for layer in layers]
+            slc.params = json.dumps(params)
+            session.merge(slc)
+            session.commit()
+    session.close()