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()