From 811d135070224294b8293faba0be29a6f1edb12a Mon Sep 17 00:00:00 2001
From: John Bodley <john.bodley@airbnb.com>
Date: Tue, 9 Jan 2018 17:41:50 -0800
Subject: [PATCH] [annotations] Fixing migration for annotation layers

---
 .../21e88bc06c02_annotation_migration.py      | 37 +++++++++++++------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/superset/migrations/versions/21e88bc06c02_annotation_migration.py b/superset/migrations/versions/21e88bc06c02_annotation_migration.py
index 77c64536ad598..4c7bb807aff04 100644
--- 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()